function sumpool(cluster::CuArray{Int}, X::CuArray{T},
                 c::Integer=length(Set(cluster))) where {T<:Real}
    dims = Dims(cluster, X)
    Y = CuArrays.zeros(T, dims.us_dims[1], c)
    scatter_add!(Y, X, cluster)
    Y
end

sumpool(cluster::Array{Int}, X::CuArray{T}, c::Integer=length(Set(cluster))) where {T<:Real} =
    sumpool(CuArray{Int64}(cluster), X, c)

function subpool(cluster::CuArray{Int}, X::CuArray{T},
                 c::Integer=length(Set(cluster))) where {T<:Real}
    dims = Dims(cluster, X)
    Y = CuArrays.zeros(T, dims.us_dims[1], c)
    scatter_sub!(Y, X, cluster)
    Y
end

subpool(cluster::Array{Int}, X::CuArray{T}, c::Integer=length(Set(cluster))) where {T<:Real} =
    subpool(CuArray{Int64}(cluster), X, c)

function prodpool(cluster::CuArray{Int}, X::CuArray{T},
                  c::Integer=length(Set(cluster))) where {T<:Real}
    dims = Dims(cluster, X)
    Y = CuArrays.ones(T, dims.us_dims[1], c)
    scatter_mul!(Y, X, cluster)
    Y
end

prodpool(cluster::Array{Int}, X::CuArray{T}, c::Integer=length(Set(cluster))) where {T<:Real} =
    prodpool(CuArray{Int64}(cluster), X, c)

function divpool(cluster::CuArray{Int}, X::CuArray{T},
                 c::Integer=length(Set(cluster))) where {T<:Real}
    dims = Dims(cluster, X)
    FT = floattype(T)
    Y = CuArrays.ones(FT, dims.us_dims[1], c)
    scatter_div!(Y, FT.(X), cluster)
    Y
end

divpool(cluster::Array{Int}, X::CuArray{T}, c::Integer=length(Set(cluster))) where {T<:Real} =
    divpool(CuArray{Int64}(cluster), X, c)

function maxpool(cluster::CuArray{Int}, X::CuArray{T},
                 c::Integer=length(Set(cluster))) where {T<:Real}
    dims = Dims(cluster, X)
    Y = CuArrays.fill(typemin(T), dims.us_dims[1], c)
    scatter_max!(Y, X, cluster)
    Y
end

maxpool(cluster::Array{Int}, X::CuArray{T}, c::Integer=length(Set(cluster))) where {T<:Real} =
    maxpool(CuArray{Int64}(cluster), X, c)

function minpool(cluster::CuArray{Int}, X::CuArray{T},
                 c::Integer=length(Set(cluster))) where {T<:Real}
    dims = Dims(cluster, X)
    Y = CuArrays.fill(typemax(T), dims.us_dims[1], c)
    scatter_min!(Y, X, cluster)
    Y
end

minpool(cluster::Array{Int}, X::CuArray{T}, c::Integer=length(Set(cluster))) where {T<:Real} =
    minpool(CuArray{Int64}(cluster), X, c)

function meanpool(cluster::CuArray{Int}, X::CuArray{T},
                  c::Integer=length(Set(cluster))) where {T<:Real}
    dims = Dims(cluster, X)
    FT = floattype(T)
    Y = CuArrays.zeros(FT, dims.us_dims[1], c)
    scatter_mean!(Y, FT.(X), cluster)
    Y
end

meanpool(cluster::Array{Int}, X::CuArray{T}, c::Integer=length(Set(cluster))) where {T<:Real} =
    meanpool(CuArray{Int64}(cluster), X, c)

@adjoint function prodpool(cluster::CuArray{Int}, X::CuArray{T}) where {T<:Real}
    prodpool(cluster, X), function (Δ)
        rev_cluster = gather_indices(cluster)
        ∇X = gather(zero(Δ)+Δ, cluster)
        @inbounds for ind = CartesianIndices(cluster)
            ind = Tuple(ind)
            inds = filter(x -> x != ind, rev_cluster[cluster[ind...]])
            for i = 1:size(X, 1)
                ∇X[i, ind...] *= mapreduce(j -> X[i, j...], *, inds; init=one(T))
            end
        end
        (nothing, ∇X)
    end
end

@adjoint function divpool(cluster::CuArray{Int}, X::CuArray{T}) where {T<:Real}
    divpool(cluster, X), function (Δ)
        rev_cluster = gather_indices(cluster)
        ∇X = -gather(zero(Δ)+Δ, cluster)
        ∇X ./= X.^2
        @inbounds for ind = CartesianIndices(cluster)
            ind = Tuple(ind)
            inds = filter(x -> x != ind, rev_cluster[cluster[ind...]])
            for i = 1:size(X, 1)
                ∇X[i, ind...] /= mapreduce(j -> X[i, j...], *, inds; init=one(T))
            end
        end
        (nothing, ∇X)
    end
end

@adjoint function maxpool(cluster::CuArray{Int}, X::CuArray{T}) where {T<:Real}
    max = maxpool(cluster, X)
    max, function (Δ)
       Δu = (X .== gather(max, cluster))
       Δu .*= gather(zero(Δ)+Δ, cluster)
       (nothing, Δu)
    end
end

@adjoint function minpool(cluster::CuArray{Int}, X::CuArray{T}) where {T<:Real}
    min = minpool(cluster, X)
    min, function (Δ)
       Δu = (X .== gather(min, cluster))
       Δu .*= gather(zero(Δ)+Δ, cluster)
       (nothing, Δu)
    end
end


using JsonGrinder, JSON, Test, SparseArrays, Mill
using HierarchicalUtils
import HierarchicalUtils: printtree

@testset "ExtractMultipleRepresentation" begin
	ex = MultipleRepresentation((ExtractCategorical(["Olda", "Tonda", "Milda"]),
		JsonGrinder.ExtractString()))
	e = ex("Olda")

	@test length(e.data) == 2
	@test size(e.data[1].data) == (4, 1)
	@test e.data[1].data == MaybeHotMatrix([2],4)
	@test size(e.data[2].data) == (2053, 1)
	@test findall(x->x > 0, SparseMatrixCSC(e.data[2].data)) .|> Tuple == [(206, 1), (272, 1), (624, 1), (738, 1), (1536, 1), (1676, 1)]

	e = ex(extractempty)
	@test nobs(e) == 0
	@test nobs(e[:e1]) == 0
	@test e[:e1].data isa MaybeHotMatrix{Int64,Array{Int64,1},Int64,Bool}
	@test nobs(e[:e2]) == 0
	@test e[:e2].data isa NGramMatrix{String,Array{String,1},Int64}

	ex2 = MultipleRepresentation((ExtractCategorical(["Olda", "Tonda", "Milda"]),
		JsonGrinder.ExtractString()))
	@test hash(ex) === hash(ex2)
	@test ex == ex2

	ex3 = MultipleRepresentation((ExtractCategorical(["Polda", "Tonda", "Milada"]),
		JsonGrinder.ExtractString()))
	@test hash(ex) !== hash(ex3)
	@test ex != ex3
end

@testset "show" begin
	ex = MultipleRepresentation((ExtractCategorical(["Olda", "Tonda", "Milda"]),
		JsonGrinder.ExtractString()))
	e = ex("Olda")

	@test buf_printtree(ex, trav=true) ==
	"""
	MultiRepresentation [""]
	  ├── e1: Categorical d = 4 ["E"]
	  └── e2: String ["U"]"""

	@test buf_printtree(e, trav=true) ==
	"""
	ProductNode with 1 obs [""]
	  ├── e1: ArrayNode(4×1 MaybeHotMatrix with Bool elements) with 1 obs ["E"]
	  └── e2: ArrayNode(2053×1 NGramMatrix with Int64 elements) with 1 obs ["U"]"""
end


using DataFrames
using Test

for testscen in 1:2
    valdir, scenario, use_permafrost, use_seaice = get_scenario(testscen)
    println(scenario)

    m = page_model()
    include("../src/components/SLRDamages.jl")

    slrdamages = addslrdamages(m)

    set_param!(m, :SLRDamages, :y_year, [2020, 2030, 2040, 2050, 2075, 2100, 2150, 2200, 2250, 2300])
    set_param!(m, :SLRDamages, :atl_adjustedtolerablelevelofsealevelrise, readpagedata(m, "test/validationdata/atl_adjustedtolerablelevelofsealevelrise.csv"))
    set_param!(m, :SLRDamages, :imp_actualreductionSLR, readpagedata(m, "test/validationdata/$valdir/imp_actualreductionSLR.csv"))
    set_param!(m, :SLRDamages, :y_year_0, 2015.)
    set_param!(m, :SLRDamages, :s_sealevel, readpagedata(m, "test/validationdata/$valdir/s_sealevel.csv"))
    set_param!(m, :SLRDamages, :cons_percap_consumption, readpagedata(m, "test/validationdata/$valdir/cons_percap_consumption.csv"))
    set_param!(m, :SLRDamages, :tct_per_cap_totalcostspercap, readpagedata(m, "test/validationdata/$valdir/tct_per_cap_totalcostspercap.csv"))
    set_param!(m, :SLRDamages, :act_percap_adaptationcosts, readpagedata(m, "test/validationdata/$valdir/act_percap_adaptationcosts.csv"))
    set_param!(m, :SLRDamages, :isatg_impactfxnsaturation, 28.333333333333336)

    ##running Model
    run(m)

    cons_percap_aftercosts = m[:SLRDamages, :cons_percap_aftercosts]
    cons_percap_aftercosts_compare = readpagedata(m, "test/validationdata/$valdir/cons_percap_aftercosts.csv")
    @test ones(10, 8) ≈ cons_percap_aftercosts ./ cons_percap_aftercosts_compare rtol=.001

    i_regionalimpactSLR = m[:SLRDamages, :i_regionalimpactSLR]
    i_regionalimpactSLR_compare = readpagedata(m, "test/validationdata/$valdir/i_regionalimpact_SLRise.csv")
    @test i_regionalimpactSLR ≈ i_regionalimpactSLR_compare rtol=.001

    iref_ImpactatReferenceGDPperCapSLR = m[:SLRDamages, :iref_ImpactatReferenceGDPperCapSLR]
    iref_ImpactatReferenceGDPperCapSLR_compare = readpagedata(m, "test/validationdata/$valdir/iref_ImpactatReferenceGDPperCap_sea.csv")
    @test iref_ImpactatReferenceGDPperCapSLR ≈ iref_ImpactatReferenceGDPperCapSLR_compare rtol=.0001

    rcons_per_cap_SLRRemainConsumption = m[:SLRDamages, :rcons_per_cap_SLRRemainConsumption]
    rcons_per_cap_SLRRemainConsumption_compare = readpagedata(m, "test/validationdata/$valdir/rcons_per_cap_SLRRemainConsumption.csv")
    @test rcons_per_cap_SLRRemainConsumption ≈ rcons_per_cap_SLRRemainConsumption_compare rtol=.01

    rgdp_per_cap_SLRRemainGDP = m[:SLRDamages, :rgdp_per_cap_SLRRemainGDP]
    rgdp_per_cap_SLRRemainGDP_compare = readpagedata(m, "test/validationdata/$valdir/rgdp_per_cap_SLRRemainGDP.csv")
    @test rgdp_per_cap_SLRRemainGDP ≈ rgdp_per_cap_SLRRemainGDP_compare rtol=.01

    igdp_ImpactatActualGDPperCapSLR = m[:SLRDamages, :igdp_ImpactatActualGDPperCapSLR]
    igdp_ImpactatActualGDPperCapSLR_compare = readpagedata(m, "test/validationdata/$valdir/igdp_ImpactatActualGDPperCap_sea.csv")
    @test igdp_ImpactatActualGDPperCapSLR ≈ igdp_ImpactatActualGDPperCapSLR_compare rtol=.01

    isat_ImpactinclSaturationandAdaptationSLR = m[:SLRDamages, :isat_ImpactinclSaturationandAdaptationSLR]
    isat_ImpactinclSaturationandAdaptationSLR_compare = readpagedata(m, "test/validationdata/$valdir/isat_ImpactinclSaturationandAdaptation_SLRise.csv")
    @test isat_ImpactinclSaturationandAdaptationSLR ≈ isat_ImpactinclSaturationandAdaptationSLR_compare rtol=.001
end


##ORTHOGONALITY HELPER
function consolidate_srcs(con_idxs::Dict{Integer,Vector{Integer}}, m::ICA_PWM_Model, obs_array::AbstractMatrix{<:Integer}, obs_lengths::AbstractVector{<:Integer}, bg_scores::AbstractMatrix{<:AbstractFloat}, contour::AbstractFloat,  models::AbstractVector{<:Model_Record}; iterates::Integer=length(m.sources)*2, remote=false)
    new_log_Li=-Inf;  iterate = 1
    T,O = size(obs_array); T=T-1; S = length(m.sources)
    new_sources=deepcopy(m.sources); new_mix=deepcopy(m.mix_matrix)
    a, cache = IPM_likelihood(new_sources, obs_array, obs_lengths, bg_scores, new_mix, true, true)

    while new_log_Li <= contour && iterate <= iterates #until we produce a model more likely than the lh contour or exceed iterates
        new_sources=deepcopy(m.sources); new_mix=deepcopy(m.mix_matrix)
        clean=Vector{Bool}(trues(O))

        for host_src in filter(!in(vcat(values(con_idxs)...)), keys(con_idxs)) #copy mix information to the source to be consolidated on as host
            for cons_src in con_idxs[host_src]
                new_mix[:,host_src].=[new_mix[o,host_src] || new_mix[o,cons_src] for o in 1:size(new_mix,1)]
            end
        end

        remote ? (merger_m = deserialize(rand(models).path)) : (merger_m = remotecall_fetch(deserialize, 1, rand(models).path)) #randomly select a model to merge
        used_m_srcs=Vector{Int64}()

        for src in unique(vcat(values(con_idxs)...)) #replace all non-host sources with sources from a merger model unlike the one being removed
            distvec=[pwm_distance(m.sources[src][1],m_src[1]) for m_src in merger_m.sources]
            m_src=findmax(distvec)[2]
            while m_src in used_m_srcs
                distvec[m_src]=0.; m_src=findmax(distvec)[2]
                length(used_m_srcs)==length(merger_m.sources) && break;break
            end

            clean[new_mix[:,src]].=false #mark dirty any obs that start with the source
            new_sources[src]=merger_m.sources[m_src]
            new_mix[:,src]=merger_m.mix_matrix[:,m_src]
            clean[new_mix[:,src]].=false #mark dirty any obs that end with the source
            push!(used_m_srcs, m_src)
        end

        if consolidate_check(new_sources)[1] #if the new sources pass the consolidate check
            new_log_Li, cache = IPM_likelihood(new_sources, obs_array, obs_lengths, bg_scores, new_mix, true, true, cache, clean) #assess likelihood
        end

        iterate += 1
    end

    return ICA_PWM_Model("candidate","consolidated $(m.origin)", new_sources, m.source_length_limits, new_mix, new_log_Li)
end

function consolidate_check(sources::AbstractVector{<:Tuple{<:AbstractMatrix{<:AbstractFloat},<:Integer}}; thresh=CONSOLIDATE_THRESH, revcomp=REVCOMP)
    pass=true
    lengthδmat=[size(src1[1],1) - size(src2[1],1) for src1 in sources, src2 in sources]
    cons_idxs=Dict{Integer,Vector{Integer}}()
    for src1 in 1:length(sources), src2 in src1+1:length(sources)
        if lengthδmat[src1,src2]==0
            revcomp ? (info_condition = (pwm_distance(sources[src1][1],sources[src2][1]) < thresh || pwm_distance(sources[src1][1],revcomp_pwm(sources[src2][1])) < thresh)) : (info_condition = (pwm_distance(sources[src1][1],sources[src2][1]) < thresh))
            if info_condition
                if !in(src1,keys(cons_idxs))
                    cons_idxs[src1]=[src2]; pass=false
                else
                    push!(cons_idxs[src1], src2)
                end
            end
        end
    end

    return pass, cons_idxs
end

                function pwm_distance(pwm1,pwm2)
                    minwml=min(size(pwm1,1),size(pwm2,1))
                    return sum([euclidean(exp.(pwm1[pos,:]), exp.(pwm2[pos,:])) for pos in 1:minwml])/minwml
                end


# Recreates the analysis from the original data files

datapath = "dataset"
projectname = "MEG3"
experimentdirnames = ["p53-Mdm2", "p53-MEG3"]

samplenames = ["A", "B", "C", "D"]

nreplicates = 3
nsamples = 4
ncells = 10

mc_iterations = 10000

outputdir = "output"
mkpath(outputdir)
outputdatapath = joinpath(outputdir, "results.jld2")

using Distributed
currentworkers = addprocs(exeflags = "--project")
@everywhere using SMLMAssociationAnalysis_NCB
using Printf
using LocalizationMicroscopy
using Statistics
using FileIO

experimentresults = Vector{Vector{Vector{Result}}}[]
for experimentdirname ∈ experimentdirnames
    println("Starting experiment $experimentdirname.")
    experimentpath = joinpath(datapath, experimentdirname)
    replicateresults = Vector{Vector{Result}}[]
    for i ∈ 1:nreplicates
        sampleresults = Vector{Result}[]
        println("    Starting replicate $i.")
        replicatepath = joinpath(experimentpath, "Replicate $i")
        for samplename ∈ samplenames
            results = Result[]
            println("        Starting sample $samplename.")
            for j ∈ 1:ncells
                println("            Starting cell $j.")
                cellpath = joinpath(replicatepath, "$samplename $(Printf.@sprintf("%03i", j)).bin.txt")
                localizations = loadlocalizations(cellpath, LocalizationMicroscopy.nikonelementstext)
                # account for variances in data collection
                if experimentdirname == experimentdirnames[2] && samplename ∈ samplenames[3:4]
                    ch1_name = "561"
                else
                    ch1_name = "647"
                end
                ch2_name = "488"
                if experimentdirname == experimentdirnames[2] &&
                   ((i == 3 && samplename ∈ samplenames[1:3]) || (i == 2 && samplename == samplenames[2]))
                    ch1_startframe = 1
                    ch2_startframe = 15001
                else
                    ch1_startframe = 1
                    ch2_startframe = 11001
                end

                ch1_task = remotecall(getmolecules, currentworkers[1], localizations,
                    ch1_name,
                    ch1_startframe,
                    11000,
                    100,
                    10,
                    34.2,
                    500,
                    200,
                )
                ch2_task = remotecall(getmolecules, currentworkers[2], localizations,
                    ch2_name,
                    ch2_startframe,
                    11000,
                    100,
                    10,
                    34.2,
                    500,
                    200,
                )

                ch1_molecules, ch1_localizations = fetch(ch1_task)
                ch2_molecules, ch2_localizations = fetch(ch2_task)

                ch1_neighbors, ch2_neighbors, distances = exclusivenearestneighbors(ch1_molecules, ch2_molecules)

                percentileranks = montecarloaffinity(
                    ch1_molecules,
                    ch2_molecules,
                    ch1_neighbors,
                    ch2_neighbors,
                    distances,
                    800,
                    mc_iterations
                )

                positivecontrol_percentileranks, positivecontrol_distances = simulate100(ch1_molecules, ch2_molecules, ch1_neighbors, ch2_neighbors, 80, 800, mc_iterations)
                # I found that the variance in running the negative control drops smoothly as a function of neighbor counts.
                # This heuristic cuts down on computational time without sacrificing quality of the correction.
                negative_iterations = length(distances) < 250 ? 30 :
                                      length(distances) < 500 ? 20 :
                                      length(distances) < 1000 ? 10 :
                                      length(distances) < 2000 ? 5 :
                                      1
                negativecontrol = [simulate0(ch1_molecules, ch2_molecules, 800, 10000) for i in 1:negative_iterations]
                negativecontrol_percentileranks = [first(i) for i in negativecontrol]
                negativecontrol_distances = [last(i) for i in negativecontrol]

                mediandistance = length(distances) == 0 ? NaN : median(distances)

                println("                Done: $(length(distances)) neighbors from $(length(ch1_molecules)) and $(length(ch2_molecules)) molecules, $(length(ch1_localizations)) and $(length(ch2_localizations)) localizations; median distance $mediandistance")
                ch1_data = ChannelData(ch1_name, ch1_molecules, ch1_neighbors)
                ch2_data = ChannelData(ch2_name, ch2_molecules, ch2_neighbors)
                result = Result(projectname, experimentdirname, 1, cellpath, i,
                                [ch1_data, ch2_data], distances, mediandistance, percentileranks,
                                positivecontrol_distances, negativecontrol_distances,
                                positivecontrol_percentileranks, negativecontrol_percentileranks)

                push!(results, result)
            end
            push!(sampleresults, results)
        end
        push!(replicateresults, sampleresults)
    end
    push!(experimentresults, replicateresults)
end

save(outputdatapath, "experimentresults", experimentresults)

rmprocs(currentworkers)


"""
```@setup general
data_dir = joinpath("..", "..", "test", "data") # hide
using CSV, DataFramesMeta, Dates, AbnormalReturns

df_firm = CSV.File(joinpath(data_dir, "daily_ret.csv")) |> DataFrame
df_mkt = CSV.File(joinpath(data_dir, "mkt_ret.csv")) |> DataFrame
df_mkt[!, :mkt] = df_mkt.mktrf .+ df_mkt.rf
df_events = CSV.File(joinpath(data_dir, "firm_earnings_announcements.csv")) |> DataFrame
mkt_data = MarketData(
    df_mkt,
    df_firm
)
```
"""
abstract type CalendarData end

"""
    struct DataVector <: CalendarData
        data::Vector{Float64}
        missing_bdays::SparseVector{Bool, Int}
        dates::ClosedInterval{Date}
        calendar::MarketCalendar
        function DataVector(data, missing_bdays, dates, calendar)
            @assert length(data) == length(missing_bdays) == bdayscount(calendar, dates.left, dates.right) + 1 "Data does not match length of dates or missings"
            new(data, missing_bdays, dates, calendar)
        end
    end

    DataVector(data::AbstractVector, dates::ClosedInterval{Date}, cal::MarketCalendar)

    DataVector(data::AbstractVector, dates::AbstractVector{Date}, cal::MarketCalendar)
"""
struct DataVector <: CalendarData
    data::Vector{Float64}
    missing_bdays::SparseVector{Bool, Int}
    dates::ClosedInterval{Date}
    calendar::MarketCalendar
    function DataVector(data, missing_bdays, dates, calendar)
        @assert length(data) == length(missing_bdays) == bdayscount(calendar, dates.left, dates.right) + 1 "Data does not match length of dates or missings"
        new(data, missing_bdays, dates, calendar)
    end
end

"""
    struct DataMatrix <: CalendarData
        data::Matrix{Float64}
        missing_bdays::SparseVector{Bool, Int}# corresponds to each row with a missing value
        dates::ClosedInterval{Date}
        calendar::MarketCalendar
        function DataMatrix(data, missing_bdays, dates, calendar)
            @assert size(data, 1) == length(missing_bdays) == bdayscount(calendar, dates.left, dates.right) + 1 "Data does not match length of dates or missings"
            new(data, missing_bdays, dates, calendar)
        end
    end

    DataMatrix(data::AbstractMatrix, dates::ClosedInterval{Date}, cal::MarketCalendar)
"""
struct DataMatrix <: CalendarData
    data::Matrix{Float64}
    missing_bdays::SparseVector{Bool, Int}# corresponds to each row with a missing value
    dates::ClosedInterval{Date}
    calendar::MarketCalendar
    function DataMatrix(data, missing_bdays, dates, calendar)
        @assert size(data, 1) == length(missing_bdays) == bdayscount(calendar, dates.left, dates.right) + 1 "Data does not match length of dates or missings"
        new(data, missing_bdays, dates, calendar)
    end
end

struct MarketData{T, MNames, FNames, N1, N2}
    calendar::MarketCalendar
    marketdata::NamedTuple{MNames, NTuple{N1, DataVector}} # column names as symbols
    firmdata::Dict{T, NamedTuple{FNames, NTuple{N2, DataVector}}} # data stored by firm id and then by column name as symbol
end

struct AllowMissing{mssng} end

"""
    mutable struct TimelineTable{Mssng, T, MNames, FNames, N1, N2} <: Tables.AbstractColumns
        parent::MarketData{T, MNames, FNames, N1, N2}
        "The current ID"
        id::T
        "Whether the functions return Vector{Union{Missing, Float64}} or Vector{Float64}"
        allow_missing::Type{AllowMissing{Mssng}}
        "Actual returnable dates that guarentees a square matrix made of the minimum and maximum of the DataVectors for this ID and set of columns"
        dates::ClosedInterval{Date}
        "Index of column names and any lags/leads"
        cols::DictIndex
        "missing days between the dates"
        missing_bdays::SparseVector{Bool, Int}
        "dates requested, which provides what data is automatically returned by functions"
        req_dates::ClosedInterval{Date}
    end

This type provides [Tables.jl](https://github.com/JuliaData/Tables.jl) access to
`DataVector`s.

Functions to update some values are:
- `AbnormalReturns.update_id!`: changes the `id` (which updates `dates` and `missing_bdays`)
- `select!`: changes the `cols` (also updates `dates` and `missing_bdays` if `cols` is different)
- `AbnormalReturns.update_dates!`: changes `req_dates` (does not update anything else)
"""
mutable struct TimelineTable{Mssng, T, MNames, FNames, N1, N2} <: Tables.AbstractColumns
    parent::MarketData{T, MNames, FNames, N1, N2}
    "The current ID"
    id::T
    "Whether the functions return Vector{Union{Missing, Float64}} or Vector{Float64}"
    allow_missing::Type{AllowMissing{Mssng}}
    "Actual returnable dates that guarentees a square matrix made of the minimum and maximum of the DataVectors for this ID and set of columns"
    dates::ClosedInterval{Date}
    "Index of column names and any lags/leads"
    cols::DictIndex
    "missing days between the dates"
    missing_bdays::SparseVector{Bool, Int}
    "dates requested, which provides what data is automatically returned by functions"
    req_dates::ClosedInterval{Date}
end

"""
    function MarketData(
        df_market,
        df_firms;
        date_col_market=:date,
        date_col_firms=:date,
        id_col=:permno,
        valuecols_market=nothing,
        valuecols_firms=nothing
    )

## Arguments
- `df_market`: A Tables.jl compatible source that stores market data, indexed by date.
    The dates must be a unique set. The column name for the date column is specified
    by the keyword argument "date_col_market"
- `df_firms`: A Tables.jl compatible source that stores firm data. Each firm must have a
    unique set of dates. The column name for the date column is specified
    by the keyword argument "date_col_firms" and the firm ID column is specified
    by the keyword argument "id_col"
- `valuecols_market=nothing`: If left as nothing, all other columns in `df_market` are
    used as the value columns. These are the columns that are stored in the resulting
    dataset. Otherwise a vector of Symbol or String specifying column names.
- `valuecols_firms=nothing`: Same as above
- `id_col=:permno`: The column corresponding to the set of firm IDs in `df_firms`

MarketData is the main data storage structure. Data is stored for each firm in
a Dict, where the data itself is a NamedTuple (names corresponding to column names,
such as "ret"), and the keys for the Dict corresponding to firm IDs. The MarketData
struct also stores overall market data and a calendar of dates.

Any firm data must have a corresponding market data date, so there cannot be a
firm return if there is not a market return on that date.

## Example

```@example general
df_firm = CSV.File(joinpath(data_dir, "daily_ret.csv"))
df_mkt = CSV.File(joinpath(data_dir, "mkt_ret.csv"))
df_mkt[!, :mkt] = df_mkt.mktrf .+ df_mkt.rf
mkt_data = MarketData(
    df_mkt,
    df_firm
)
```
"""
function MarketData(
    df_market,
    df_firms;
    date_col_market=:date,
    date_col_firms=:date,
    id_col=:permno,
    valuecols_market=nothing,
    valuecols_firms=nothing
)
    df_market = DataFrame(df_market)
    df_firms = DataFrame(df_firms)
    if valuecols_market === nothing
        valuecols_market = Symbol.([n for n in Symbol.(names(df_market)) if n ∉ [date_col_market]])
    end
    if valuecols_firms === nothing
        valuecols_firms = Symbol.([n for n in Symbol.(names(df_firms)) if n ∉ [date_col_firms, id_col]])
    end

    df_market = select(df_market, vcat([date_col_market], valuecols_market))
    dropmissing!(df_market, date_col_market)
    #dropmissing!(df_market)
    sort!(df_market)

    if any(nonunique(df_market, [date_col_market]))
        @error("There are duplicate date rows in the market data")
    end

    df_firms = select(df_firms, vcat([id_col, date_col_firms], valuecols_firms))
    dropmissing!(df_firms, [id_col, date_col_firms])
    sort!(df_firms, [id_col, date_col_firms])

    if any(nonunique(df_firms, [id_col, date_col_firms]))
        @error("There are duplicate id-date rows in the firm data")
    end

    cal = MarketCalendar(df_market[:, date_col_market])

    market_data = NamedTuple(
        valuecols_market .=>
        DataVector.(
            Tables.columns(df_market[:, valuecols_market]),
            Ref(df_market[:, date_col_market]),
            Ref(cal)
        )
    )

    check_all_businessdays(unique(df_firms[:, date_col_firms]), cal)

    gdf = groupby(df_firms, id_col)
    df_temp = combine(
        gdf,
        date_col_firms => (x -> add_missing_bdays(x, cal)) => date_col_firms
    )
    if nrow(df_temp) > 0
        insertcols!(df_temp, [col => missing for col in valuecols_firms]...)
        df_firms = vcat(
            df_firms,
            df_temp
        )
        sort!(df_firms, [id_col, date_col_firms])
        gdf = groupby(df_firms, id_col)
    end

    col_tab = columntable(gdf)

    firm_data = Dict{typeof(col_tab[id_col][1][1]), NamedTuple{Tuple(valuecols_firms), NTuple{length(valuecols_firms), DataVector}}}()
    sizehint!(firm_data, length(col_tab[id_col]))

    for i in 1:length(col_tab[id_col])
        firm_data[col_tab[id_col][i][1]] = NamedTuple(
            valuecols_firms .=> (
                DataVector(
                    col_tab[col][i],
                    col_tab[date_col_firms][i],
                    cal
                ) for col in valuecols_firms
            )
        )
    end

    MarketData(
        cal,
        market_data,
        firm_data
    )
end

function check_all_businessdays(dates, cal)
    bday_list = isbday(cal, dates)
    if !all(bday_list)
        bday_list_inv = (!).(bday_list)
        if sum(bday_list_inv) <= 3
            @error("Dates $(dates[bday_list_inv]) are not in the MARKET_DATA_CACHE")
        else
            d1_test = findfirst(bday_list_inv)
            d2_test = findlast(bday_list_inv)
            @error("Dates $(dates[d1_test]) ... $(dates[d2_test]) are not in the MARKET_DATA_CACHE")
        end
    end
end

function setup_calendar_data(f, data, dates::ClosedInterval{Date}, cal::MarketCalendar)
    if any(ismissing.(data))
        missing_days = sparsevec(any(ismissing, data, dims=2))
        data = coalesce.(data, zero(nonmissingtype(eltype(data))))
    else
        missing_days = spzeros(Bool, size(data, 1))
    end
    f(
        data,
        missing_days,
        dates,
        cal
    )
end

function DataVector(data::AbstractVector, dates::ClosedInterval{Date}, cal::MarketCalendar)
    setup_calendar_data(
        DataVector,
        data,
        dates,
        cal
    )
end

function DataVector(data::AbstractVector, dates::AbstractVector{Date}, cal::MarketCalendar)
    if any(ismissing.(data))
        if all(ismissing.(data))
            return DataVector(
                zeros(nonmissingtype(eltype(data)), 1),
                sparsevec([1], [true]),
                dates[1] .. dates[1],
                cal
            )
        end
        i = findfirst(!ismissing, data)
        j = findlast(!ismissing, data)
        setup_calendar_data(
            DataVector,
            data[i:j],
            dates[i] .. dates[j],
            cal
        )
    else
        setup_calendar_data(
            DataVector,
            data,
            dates[1] .. dates[end],
            cal
        )
    end
end

function DataMatrix(data::AbstractMatrix, dates::ClosedInterval{Date}, cal::MarketCalendar)
    setup_calendar_data(
        DataMatrix,
        data,
        dates,
        cal
    )
end
function add_missing_bdays(dates, cal)
    out = Date[]
    if length(dates) == bdayscount(cal, dates[1], dates[end]) + 1
        return out
    end
    dates_counter = 1
    for d in listbdays(cal, dates[1], dates[end])
        if dates[dates_counter] > d
            push!(out, d)
        else
            dates_counter += 1
        end
    end
    out
end

function get_missing_bdays(cal::MarketCalendar, missing_bdays::SparseVector{Bool, Int}, dates_missings::ClosedInterval{Date}, new_dates::ClosedInterval{Date})
    missing_bdays[date_range(cal, dates_missings, new_dates)]
end

function get_missing_bdays(data, new_dates::ClosedInterval{Date})
    get_missing_bdays(calendar(data), data_missing_bdays(data), data_dates(data), new_dates)
end

function date_range(cal::MarketCalendar, timeline_dates::ClosedInterval{Date}, new_dates::ClosedInterval{Date})
    dates = dates_min_max(timeline_dates, new_dates)
    s = bdayscount(cal, dt_min(timeline_dates), dt_min(dates)) + 1
    e = s + bdayscount(cal, dt_min(dates), dt_max(dates)) - !isbday(cal, dt_max(dates))
    s:e
end

function date_range(cal::MarketCalendar, data::CalendarData, dt_min::Date, dt_max::Date)
    date_range(cal, data_dates(data), dt_min .. dt_max)
end

function date_range(cal::MarketCalendar, data::CalendarData, dates::ClosedInterval{Date})
    date_range(cal, data_dates(data), dates)
end

# , mssngs::SparseVector{Bool, Int}=spzeros(Bool, size(raw_values(data), 1))
function Base.getindex(data::DataMatrix, dates::ClosedInterval{Date})
    new_dates = dates_min_max(data_dates(data), dates)
    r = date_range(calendar(data), data_dates(data), new_dates)
    data.data[r, :]
end

function Base.getindex(data::DataVector, dates::ClosedInterval{Date})
    new_dates = dates_min_max(data_dates(data), dates)
    r = date_range(calendar(data), data_dates(data), new_dates)
    data.data[r]
end

function combine_missing_bdays(vals::SparseVector{Bool, Int}...)::SparseVector{Bool, Int}
    (|).(vals...)
end

function combine_missing_bdays(vals::Vector{SparseVector{Bool, Int}})::SparseVector{Bool, Int}
    out = sparsevec(zeros(Bool, length(vals[1])))
    for v in vals
        out = out .| v
    end
    out
end

function Base.view(data::DataVector, dates::ClosedInterval{Date}, mssngs::SparseVector{Bool, Int})
    @assert dt_min(data) <= dt_min(dates) && dt_max(dates) <= dt_max(data) "Date range out of bounds"
    r = date_range(calendar(data), data_dates(data), dates)
    @assert length(mssngs) == length(r) "Missing data vector is the wrong length"
    if nnz(mssngs) == 0
        view(data.data, r)
    else
        view(data.data, r[(!).(mssngs)])
    end
end

function Base.view(data::DataMatrix, dates::ClosedInterval{Date}, mssngs::SparseVector{Bool, Int})
    @assert dt_min(data) <= dt_min(dates) && dt_max(dates) <= dt_max(data) "Date range out of bounds"
    r = date_range(calendar(data), data_dates(data), dates)
    @assert length(mssngs) == length(r) "Missing data vector is the wrong length"
    if nnz(mssngs) == 0
        view(data.data, r, :)
    else
        view(data.data, r[(!).(mssngs)], :)
    end
end

function Base.view(data::DataVector, dates::ClosedInterval{Date})
    @assert dt_min(data) <= dt_min(dates) && dt_max(dates) <= dt_max(data) "Date range out of bounds"
    r = date_range(calendar(data), data_dates(data), dates)
    view(data.data, r)
end

function Base.view(data::DataMatrix, dates::ClosedInterval{Date})
    @assert dt_min(data) <= dt_min(dates) && dt_max(dates) <= dt_max(data) "Date range out of bounds"
    r = date_range(calendar(data), data_dates(data), dates)
    view(data.data, r, :)
end

Base.getindex(data::NamedTuple{Names, NTuple{N, DataVector}}, col::TimelineColumn) where {Names, N} = shift(data[Symbol(col)], col.shifts)

function Base.getindex(data::MarketData{T, MNames, FNames}, id::T, col::TimelineColumn) where {T, MNames, FNames}
    shift(pick_data(data, id, Symbol(col)), col.shifts)
end

function Base.getindex(
    data::MarketData{T, MNames, FNames},
    id::T,
    cols::Vector{TimelineColumn}=TimelineColumn.([FNames..., MNames...]),
    allow_mssng::Type{AllowMissing{Mssng}}=AllowMissing{false},
) where {T, MNames, FNames, Mssng}
    real_dates = dates_min_max(get_all_dates(data, id, cols))
    data[id, real_dates, cols, allow_mssng]
end

function Base.getindex(
    data::MarketData{T, MNames, FNames},
    id::T,
    dates::ClosedInterval{Date},
    cols::Vector{TimelineColumn}=TimelineColumn.([FNames..., MNames...]),
    allow_mssng::Type{AllowMissing{Mssng}}=AllowMissing{false}
) where {T, MNames, FNames, Mssng}
    out = TimelineTable(
        data,
        id,
        allow_mssng,
        dates_min_max(get_all_dates(data, id, cols)),
        DictIndex(cols),
        sparsevec(zeros(Bool, bdayscount(calendar(data), dt_min(dates), dt_max(dates)))),
        dates,
    )
    out.missing_bdays = combine_all_missing_bdays(out)
    out
end

function Base.getindex(
    data::MarketData{T},
    id::T,
    dates::ClosedInterval{Date},
    cols::Vector,
    allow_mssng::Type{AllowMissing{Mssng}}=AllowMissing{false}
) where {T, Mssng}
    data[id, dates, TimelineColumn.(cols), allow_mssng]
end

function Base.getindex(data::TimelineTable, col::TimelineColumn)
    parent(data)[data.id, col]
end
Base.getindex(data::TimelineTable, col::Symbol) = data[TimelineColumn(col)]

function Base.getindex(
    data::TimelineTable,
    ::Colon,
    col::TimelineColumn
)
    ret_dates = dates_min_max(data_dates(data), norm_dates(data))
    data[ret_dates, col]
end

function Base.getindex(
    data::TimelineTable{false},
    dates::ClosedInterval{Date},
    col::TimelineColumn
)
    new_mssngs = get_missing_bdays(calendar(data), data_missing_bdays(data), data_dates(data), dates)
    if nnz(new_mssngs) == 0
        data[col][dates]
    else
        data[col][dates][(!).(new_mssngs)]
    end
end

function Base.getindex(
    data::TimelineTable{true},
    dates::ClosedInterval{Date},
    col::TimelineColumn
)
    new_mssngs = get_missing_bdays(calendar(data), data_missing_bdays(data), data_dates(data), dates)
    out = data[col][dates]
    out = allowmissing(out)
    out[new_mssngs] .= missing
    to_add_pre = ifelse(
        dt_min(data) < dt_min(data_dates(data)),
        bdayscount(calendar(data), dt_min(data), dt_min(data_dates(data))),
        0
    )
    to_add_post = ifelse(
        dt_max(data) > dt_max(data_dates(data)),
        bdayscount(calendar(data), dt_max(data_dates(data)), dt_max(data)),
        0
    )
    vcat(
        fill(missing, to_add_pre),
        out,
        fill(missing, to_add_post)
    )
end

function Base.getindex(
    data::TimelineTable,
    ::Colon,
    col::Symbol
)
    data[:, TimelineColumn(col)]
end

function Base.getindex(
    data::TimelineTable,
    dates::ClosedInterval{Date},
    col::Symbol
)
    data[dates, TimelineColumn(col)]
end

function Base.getproperty(
    data::TimelineTable,
    sym::Symbol
)
    if sym == :calendar
        getfield(getfield(data, :parent), :calendar)
    elseif sym == :firmdata
        getfield(getfield(data, :parent), :firmdata)
    elseif sym == :marketdata
        getfield(getfield(data, :parent), :marketdata)
    else
        getfield(data, sym)
    end
end

function get_dates(data::TimelineTable{true})
    listbdays(calendar(data), dt_min(data), dt_max(data))
end

function get_dates(data::TimelineTable{false})
    real_dates = dates_min_max(data_dates(data), norm_dates(data))
    out = listbdays(calendar(data), dt_min(real_dates), dt_max(real_dates))
    new_mssngs = get_missing_bdays(calendar(data), data_missing_bdays(data), data_dates(data), real_dates)
    if nnz(new_mssngs) == 0
        return out
    else
        return out[(!).(new_mssngs)]
    end
end

Base.names(x::TimelineTable) = x.cols.cols

Tables.istable(::Type{<:TimelineTable}) = true
Tables.columnaccess(::Type{<:TimelineTable}) = true
function Tables.schema(obj::TimelineTable{Mssngs}) where {Mssngs}
    col_sym = Tables.columnnames(obj)
    col_types = vcat(
        [Date],
        fill(Mssngs ? Union{Float64, Missing} : Float64, length(col_sym)-1)
    )
    Tables.Schema(col_sym, col_types)
end

Tables.columns(x::TimelineTable) = x
function Tables.getcolumn(x::TimelineTable, i::Int)
    if i == 1
        get_dates(x)
    else
        Tables.getcolumn(x, x.cols.lookup[i-1])# subtract 1 since date is the first column produced
    end
end
function Tables.getcolumn(x::TimelineTable, nm::TimelineColumn)
    x[:, nm]
end
function Tables.getcolumn(x::TimelineTable, nm::Symbol)
    if nm == :date
        return get_dates(x)
    end
    if nm ∈ Tables.columnnames(x)
        pos = findfirst(nm .== Tables.columnnames(x))
        return Tables.getcolumn(x, pos)
    end
end

function Tables.columnnames(x::TimelineTable)
    vcat([:date], Symbol.(String.(names(x))))
end

function Base.length(data::TimelineTable{false})
    real_dates = dates_min_max(data_dates(data), norm_dates(data))
    c = bdayscount(calendar(data), dt_min(real_dates), dt_max(real_dates)) + isbday(calendar(data), dt_max(real_dates))
    new_mssngs = get_missing_bdays(calendar(data), data_missing_bdays(data), data_dates(data), real_dates)
    return c - nnz(new_mssngs)
end

function Base.length(data::TimelineTable{true})
    bdayscount(calendar(data), dt_min(data), dt_max(data)) + isbday(calendar(data), dt_max(data))
end

function Base.length(x::DataVector)
    return length(data_missing_bdays(x)) - nnz(data_missing_bdays(x))
end

DataFrames.dropmissing(data::TimelineTable{false}) = data
DataFrames.allowmissing(data::TimelineTable{true}) = data

function DataFrames.dropmissing(data::TimelineTable{true})
    TimelineTable(
        parent(data),
        data_id(data),
        AllowMissing{false},
        data_dates(data),
        DictIndex(names(data)),
        data_missing_bdays(data),
        norm_dates(data)
    )
end

function DataFrames.allowmissing(data::TimelineTable{false})
    TimelineTable(
        parent(data),
        data_id(data),
        AllowMissing{true},
        data_dates(data),
        DictIndex(names(data)),
        data_missing_bdays(data),
        norm_dates(data)
    )
end

function update_dates!(
    data::TimelineTable,
    dates::ClosedInterval{Date}
)
    data.req_dates = dates
    data
end

function update_id!(
    data::TimelineTable{Mssng, T},
    id::T
) where {Mssng, T}
    if id != data.id
        data.id = id
        data.dates = dates_min_max(get_all_dates(data))
        data.missing_bdays = combine_all_missing_bdays(data)
    end
    data
end

function DataFrames.select!(x::TimelineTable, cols::Vector)
    select!(x, TimelineColumn.(cols))
end
function DataFrames.select!(data::TimelineTable, cols::Vector{TimelineColumn})
    updates = sort(cols) != sort(names(data)) # if the columns are changing, then update dates and missing_bdays
    data.cols = DictIndex(cols)
    if updates
        data.dates = dates_min_max(get_all_dates(data))
        data.missing_bdays = combine_all_missing_bdays(data)
    end
    data
end

function get_all_dates(data::MarketData{T}, id::T, cols::Vector{TimelineColumn}) where {T}
    shift_dates.(
        pick_data.(Ref(data), id, Symbol.(cols)),
        shift_count.(cols)
    )
end

function get_all_dates(data::TimelineTable)
    get_all_dates(parent(data), data.id, names(data))
end

function combine_all_missing_bdays(data::TimelineTable)
    out = sparsevec(
        zeros(
            Bool,
            bdayscount(
                calendar(data),
                dt_min(data_dates(data)),
                dt_max(data_dates(data))
            ) + 1
        )
    )
    for col in names(data)
        mssngs = data_missing_bdays(data[Symbol(col)]) # get based on symbol here since if there are no missings
        # then skip the next step, if there are missings then fetch the new vector and adjust
        if nnz(mssngs) > 0
            out = out .| get_missing_bdays(data[col], data_dates(data))
        end
    end
    out
end

raw_values(x::CalendarData) = x.data
data_dates(x::CalendarData) = x.dates
data_missing_bdays(x::CalendarData) = x.missing_bdays

data_dates(x::TimelineTable) = x.dates
data_missing_bdays(x::TimelineTable) = x.missing_bdays
norm_dates(x::TimelineTable) = x.req_dates
parent(data::TimelineTable) = data.parent
data_id(data::TimelineTable) = data.id

calendar(x::CalendarData) = x.calendar
calendar(x::MarketData) = x.calendar
calendar(x::TimelineTable) = calendar(parent(x))

firmdata(data::MarketData{T}, i::T, sym::Symbol) where {T} = data.firmdata[i][sym]
marketdata(data::MarketData, sym::Symbol) = data.marketdata[sym]
pick_data(data::MarketData{T, MNames, FNames}, id::T, sym::Symbol) where {T, MNames, FNames} = sym ∈ FNames ? firmdata(data, id, sym) : marketdata(data, sym)

pick_data(data::MarketData{T}, id::T, col::TimelineColumn) where {T} = pick_data(data, id, Symbol(col))

data_dates(x::MarketData) = x.dates

data_loc(data::MarketData{T, MNames, FNames}, sym::Symbol) where {T, MNames, FNames} = sym ∈ FNames ? :firmdata : :marketdata
function data_loc(data::MarketData{T, MNames, FNames}, syms::Vector{Symbol}) where {T, MNames, FNames}
    all_locs = data_loc.(Ref(data), syms)
    (
        syms[all_locs .== :firmdata],
        syms[all_locs .== :marketdata]
    )
end

function Base.show(io::IO, data::MarketData{T, MNames, FNames}) where {T, MNames, FNames}
    println(io, "MarketData with ID type $T with $(length(getfield(data, :firmdata))) unique firms")
    println(io, data.calendar)
    println(io, "Market Columns: $(join(MNames, ", "))")
    println(io, "Firm Columns: $(join(FNames, ", "))")
end

function Base.show(io::IO, data::TimelineTable{Mssng, T, MNames, FNames}) where {Mssng, T, MNames, FNames}
    println(io, "TimelineTable for $(data_id(data)) and available columns $(join([MNames..., FNames...], ","))")
    println(io, "Maximum dates given current column selection: $(data_dates(data))")
    println(io, "Currently selected dates: $(norm_dates(data))")
    pretty_table(io, data; header = vcat(["Date"], names(data)), tf=tf_simple)
end

dt_min(x::ClosedInterval{Date}) = x.left
dt_min(x::TimelineTable) = dt_min(norm_dates(x))
dt_min(x::CalendarData) = dt_min(data_dates(x))
dt_max(x::ClosedInterval{Date}) = x.right
dt_max(x::TimelineTable) = dt_max(norm_dates(x))
dt_max(x::CalendarData) = dt_max(data_dates(x))

cal_dt_min(x::CalendarData) = cal_dt_min(calendar(x))
cal_dt_max(x::CalendarData) = cal_dt_max(calendar(x))
cal_dt_min(x::TimelineTable) = cal_dt_min(calendar(x))
cal_dt_max(x::TimelineTable) = cal_dt_max(calendar(x))
cal_dt_min(x::MarketData) = cal_dt_min(calendar(x))
cal_dt_max(x::MarketData) = cal_dt_max(calendar(x))


using MOCNeutronTransport

# Model setup
# ----------------------------------------------------------------------------------------
gmsh.initialize()
gmsh.option.set_number("Geometry.OCCParallel", 1)
gmsh.option.set_number("General.NumThreads", 32)

# Import CAD file
gmsh.merge("/home/kcvaughn/Downloads/Core Assembly - Test Layout 2 - 6-30-20 Export.STEP")

# Shift the model to the origin of the xy-plane
dim_tags = gmsh.model.get_entities(3)
# Rotate the entities in dim_tags in the built-in CAD representation by angle radians
# around the axis of revolution defined by the point (x, y, z) and the direction (ax, ay, az).
# gmsh.model.occ.rotate(dim_tags, x, y, z, ax, ay, az, angle)
gmsh.model.occ.rotate(dim_tags, 0, 0, 0, 1, 0, 0, π/2)
# Translate it so that the corner is at the origin
#gmsh.model.occ.translate(dim_tags, 37.3225 - 0.0099 + 10*grid_offset,
#                                   44.09 - 0.00174209 + 10*grid_offset,
#                                   -436.5625)
# Model is in mm. Convert to cm by shrinking everything by 1/10
#gmsh.model.occ.dilate(dim_tags, 0, 0, 0, 1//10, 1//10, 0)

rect_tag = gmsh.model.occ.add_rectangle(-500.0, -500.0, 300, 1000.0, 1000.0)

# Synchronize the CAD and Gmsh models
gmsh.model.occ.synchronize()

#Geometry.OCCParallel
# Verify results
# gmsh.fltk.run()
intersections = gmsh.model.occ.intersect([(Int32(2), rect_tag)], dim_tags)
gmsh.model.occ.synchronize()
gmsh.write("FNR_slice.step")
#gmsh.finalize()


# wrappers of low-level functionality

using CUDAapi: cudaDataType
using CUDAdrv: CuDefaultStream, CuStream

function version()
    ver = cutensorGetVersion()
    major, ver = divrem(ver, 10000)
    minor, patch = divrem(ver, 100)

    VersionNumber(major, minor, patch)
end

function cuda_version()
  ver = cutensorGetCudartVersion()
  major, ver = divrem(ver, 1000)
  minor, patch = divrem(ver, 10)

  VersionNumber(major, minor, patch)
end

const ModeType = AbstractVector{<:Union{Char, Integer}}

is_unary(op::cutensorOperator_t) =
    (op ∈ (CUTENSOR_OP_IDENTITY, CUTENSOR_OP_SQRT, CUTENSOR_OP_RELU, CUTENSOR_OP_CONJ,
            CUTENSOR_OP_RCP))
is_binary(op::cutensorOperator_t) =
    (op ∈ (CUTENSOR_OP_ADD, CUTENSOR_OP_MUL, CUTENSOR_OP_MAX, CUTENSOR_OP_MIN))

mutable struct CuTensorDescriptor
    desc::Ref{cutensorTensorDescriptor_t}

    function CuTensorDescriptor(a; size = size(a), strides = strides(a), eltype = eltype(a),
                                   op = CUTENSOR_OP_IDENTITY)
        sz = collect(Int64, size)
        st = collect(Int64, strides)
        desc = Ref{cutensorTensorDescriptor_t}()
        cutensorInitTensorDescriptor(handle(), desc, length(sz), sz, st,
                                     cudaDataType(eltype), op)
        obj = new(desc)
        return obj
    end
end

Base.cconvert(::Type{Ptr{cutensorTensorDescriptor_t}}, obj::CuTensorDescriptor) = obj.desc

function elementwiseTrinary!(
    alpha::Number, A::CuArray, Ainds::ModeType, opA::cutensorOperator_t,
    beta::Number,  B::CuArray, Binds::ModeType, opB::cutensorOperator_t,
    gamma::Number, C::CuArray{T}, Cinds::ModeType, opC::cutensorOperator_t,
    D::CuArray{T}, Dinds::ModeType, opAB::cutensorOperator_t,
    opABC::cutensorOperator_t; stream::CuStream=CuDefaultStream()) where {T}

    !is_unary(opA)    && throw(ArgumentError("opA must be a unary op!"))
    !is_unary(opB)    && throw(ArgumentError("opB must be a unary op!"))
    !is_unary(opC)    && throw(ArgumentError("opC must be a unary op!"))
    !is_binary(opAB)  && throw(ArgumentError("opAB must be a binary op!"))
    !is_binary(opABC) && throw(ArgumentError("opABC must be a binary op!"))
    descA = CuTensorDescriptor(A; op = opA)
    descB = CuTensorDescriptor(B; op = opB)
    descC = CuTensorDescriptor(C; op = opC)
    @assert size(C) == size(D) && strides(C) == strides(D)
    descD = descC # must currently be identical
    #typeCompute = cudaDataType(T)
    typeCompute = cudaDataType(T)
    modeA = collect(Cint, Ainds)
    modeB = collect(Cint, Binds)
    modeC = collect(Cint, Cinds)
    modeD = modeC
    cutensorElementwiseTrinary(handle(),
                                T[alpha], A, descA, modeA,
                                T[beta],  B, descB, modeB,
                                T[gamma], C, descC, modeC,
                                          D, descD, modeD,
                                opAB, opABC, typeCompute, stream)
    return D
end

function elementwiseTrinary!(
    alpha::Number, A::Array, Ainds::ModeType, opA::cutensorOperator_t,
    beta::Number, B::Array, Binds::ModeType, opB::cutensorOperator_t,
    gamma::Number, C::Array{T}, Cinds::ModeType, opC::cutensorOperator_t,
    D::Array{T}, Dinds::ModeType, opAB::cutensorOperator_t,
    opABC::cutensorOperator_t; stream::CuStream=CuDefaultStream()) where {T}

    !is_unary(opA)    && throw(ArgumentError("opA must be a unary op!"))
    !is_unary(opB)    && throw(ArgumentError("opB must be a unary op!"))
    !is_unary(opC)    && throw(ArgumentError("opC must be a unary op!"))
    !is_binary(opAB)  && throw(ArgumentError("opAB must be a binary op!"))
    !is_binary(opABC) && throw(ArgumentError("opABC must be a binary op!"))
    descA = CuTensorDescriptor(A; op = opA)
    descB = CuTensorDescriptor(B; op = opB)
    descC = CuTensorDescriptor(C; op = opC)
    @assert size(C) == size(D) && strides(C) == strides(D)
    descD = descC # must currently be identical
    typeCompute = cudaDataType(T)
    modeA = collect(Cint, Ainds)
    modeB = collect(Cint, Binds)
    modeC = collect(Cint, Cinds)
    modeD = modeC
    cutensorElementwiseTrinary(handle(),
                               T[alpha], A, descA, modeA,
                               T[beta],  B, descB, modeB,
                               T[gamma], C, descC, modeC,
                                         D, descD, modeD,
                               opAB, opABC, typeCompute, stream)
    return D
end

function elementwiseBinary!(
    alpha::Number, A::CuArray, Ainds::ModeType, opA::cutensorOperator_t,
    gamma::Number, C::CuArray{T}, Cinds::ModeType, opC::cutensorOperator_t,
    D::CuArray{T}, Dinds::ModeType, opAC::cutensorOperator_t;
    stream::CuStream=CuDefaultStream()) where {T}

    !is_unary(opA)    && throw(ArgumentError("opA must be a unary op!"))
    !is_unary(opC)    && throw(ArgumentError("opC must be a unary op!"))
    !is_binary(opAC)  && throw(ArgumentError("opAC must be a binary op!"))
    descA = CuTensorDescriptor(A; op = opA)
    descC = CuTensorDescriptor(C; op = opC)
    @assert size(C) == size(D) && strides(C) == strides(D)
    descD = descC # must currently be identical
    typeCompute = cudaDataType(T)
    modeA = collect(Cint, Ainds)
    modeC = collect(Cint, Cinds)
    modeD = modeC
    cutensorElementwiseBinary(handle(),
                              T[alpha], A, descA, modeA,
                              T[gamma], C, descC, modeC,
                                        D, descD, modeD,
                              opAC, typeCompute, stream)
    return D
end

function elementwiseBinary!(
    alpha::Number, A::Array, Ainds::ModeType, opA::cutensorOperator_t,
    gamma::Number, C::Array{T}, Cinds::ModeType, opC::cutensorOperator_t,
    D::Array{T}, Dinds::ModeType, opAC::cutensorOperator_t;
    stream::CuStream=CuDefaultStream()) where {T}

    !is_unary(opA)    && throw(ArgumentError("opA must be a unary op!"))
    !is_unary(opC)    && throw(ArgumentError("opC must be a unary op!"))
    !is_binary(opAC)  && throw(ArgumentError("opAC must be a binary op!"))
    descA = CuTensorDescriptor(A; op = opA)
    descC = CuTensorDescriptor(C; op = opC)
    @assert size(C) == size(D) && strides(C) == strides(D)
    descD = descC # must currently be identical
    typeCompute = cudaDataType(T)
    modeA = collect(Cint, Ainds)
    modeC = collect(Cint, Cinds)
    modeD = modeC
    cutensorElementwiseBinary(handle(),
                              T[alpha], A, descA, modeA,
                              T[gamma], C, descC, modeC,
                                        D, descD, modeD,
                              opAC, typeCompute, stream)
    return D
end

function elementwiseBinary!(
    alpha::Number, A::CuTensor, opA::cutensorOperator_t, gamma::Number, C::CuTensor{T}, opC::cutensorOperator_t, D::CuTensor{T}, opAC::cutensorOperator_t; stream::CuStream=CuDefaultStream()) where {T}

    !is_unary(opA)    && throw(ArgumentError("opA must be a unary op!"))
    !is_unary(opC)    && throw(ArgumentError("opC must be a unary op!"))
    !is_binary(opAC)  && throw(ArgumentError("opAC must be a binary op!"))
    descA = CuTensorDescriptor(A; op = opA)
    descC = CuTensorDescriptor(C; op = opC)
    @assert size(C) == size(D) && strides(C) == strides(D)
    descD = descC # must currently be identical
    typeCompute = cudaDataType(T)
    cutensorElementwiseBinary(handle(),
                              T[alpha], A.data, descA, A.inds,
                              T[gamma], C.data, descC, C.inds,
                                        D.data, descD, C.inds,
                              opAC, typeCompute, stream)
    return D
end

function permutation!(alpha::Number, A::CuArray, Ainds::ModeType,
                      B::CuArray, Binds::ModeType; stream::CuStream=CuDefaultStream())
    #!is_unary(opPsi)    && throw(ArgumentError("opPsi must be a unary op!"))
    descA = CuTensorDescriptor(A)
    descB = CuTensorDescriptor(B)
    T = eltype(B)
    typeCompute = cudaDataType(T)
    modeA = collect(Cint, Ainds)
    modeB = collect(Cint, Binds)
    cutensorPermutation(handle(), T[alpha], A, descA, modeA, B, descB, modeB, typeCompute,
                        stream)
    return B
end
function permutation!(alpha::Number, A::Array, Ainds::ModeType,
                      B::Array, Binds::ModeType; stream::CuStream=CuDefaultStream())
    #!is_unary(opPsi)    && throw(ArgumentError("opPsi must be a unary op!"))
    descA = CuTensorDescriptor(A)
    descB = CuTensorDescriptor(B)
    T = eltype(B)
    typeCompute = cudaDataType(T)
    modeA = collect(Cint, Ainds)
    modeB = collect(Cint, Binds)
    cutensorPermutation(handle(), T[alpha], A, descA, modeA, B, descB, modeB, typeCompute,
                        stream)
    return B
end

function contraction!(
    alpha::Number, A::CuArray, Ainds::ModeType, opA::cutensorOperator_t,
                   B::CuArray, Binds::ModeType, opB::cutensorOperator_t,
    beta::Number,  C::CuArray, Cinds::ModeType, opC::cutensorOperator_t,
                                                opOut::cutensorOperator_t;
    pref::cutensorWorksizePreference_t=CUTENSOR_WORKSPACE_RECOMMENDED,
    algo::cutensorAlgo_t=CUTENSOR_ALGO_DEFAULT, stream::CuStream=CuDefaultStream(),
    compute_type::Type=eltype(C), plan::Union{cutensorContractionPlan_t, Nothing}=nothing)

    !is_unary(opA)    && throw(ArgumentError("opA must be a unary op!"))
    !is_unary(opB)    && throw(ArgumentError("opB must be a unary op!"))
    !is_unary(opC)    && throw(ArgumentError("opC must be a unary op!"))
    !is_unary(opOut)  && throw(ArgumentError("opOut must be a unary op!"))
    descA = CuTensorDescriptor(A; op = opA)
    descB = CuTensorDescriptor(B; op = opB)
    descC = CuTensorDescriptor(C; op = opC)
    # for now, D must be identical to C (and thus, descD must be identical to descC)
    T = eltype(C)
    computeType = cutensorComputeType(T) #CUTENSOR_R_MIN_64F #TODO cudaDataType(T)
    modeA = collect(Cint, Ainds)
    modeB = collect(Cint, Binds)
    modeC = collect(Cint, Cinds)

    alignmentRequirementA = Ref{UInt32}(C_NULL)
    cutensorGetAlignmentRequirement(handle(), A, descA, alignmentRequirementA)
    alignmentRequirementB = Ref{UInt32}(C_NULL)
    cutensorGetAlignmentRequirement(handle(), B, descB, alignmentRequirementB)
    alignmentRequirementC = Ref{UInt32}(C_NULL)
    cutensorGetAlignmentRequirement(handle(), C, descC, alignmentRequirementC)
    desc = Ref(cutensorContractionDescriptor_t(ntuple(i->0, Val(256))))
    cutensorInitContractionDescriptor(handle(),
                                      desc,
                   descA, modeA, alignmentRequirementA[],
                   descB, modeB, alignmentRequirementB[],
                   descC, modeC, alignmentRequirementC[],
                   descC, modeC, alignmentRequirementC[],
                   computeType)
    find = Ref(cutensorContractionFind_t(ntuple(i->0, Val(64))))
    cutensorInitContractionFind(handle(), find, algo)

    @workspace fallback=1<<27 size=@argout(
            cutensorContractionGetWorkspace(handle(), desc, find, pref,
                                            out(Ref{UInt64}(C_NULL)))
        )[] workspace->begin
            plan_ref = Ref(cutensorContractionPlan_t(ntuple(i->0, Val(640))))
            if isnothing(plan)
                cutensorInitContractionPlan(handle(), plan_ref, desc, find, sizeof(workspace))
            else
                plan_ref = Ref(plan)
            end
            cutensorContraction(handle(), plan_ref,
                                T[convert(T, alpha)], A, B,
                                T[convert(T, beta)],  C, C,
                                workspace, sizeof(workspace), stream)
        end
    return C
end

function plan_contraction(
    A::CuArray, Ainds::ModeType, opA::cutensorOperator_t,
    B::CuArray, Binds::ModeType, opB::cutensorOperator_t,
    C::CuArray, Cinds::ModeType, opC::cutensorOperator_t,
                                                opOut::cutensorOperator_t;
    pref::cutensorWorksizePreference_t=CUTENSOR_WORKSPACE_RECOMMENDED,
    algo::cutensorAlgo_t=CUTENSOR_ALGO_DEFAULT, compute_type::Type=eltype(C))

    !is_unary(opA)    && throw(ArgumentError("opA must be a unary op!"))
    !is_unary(opB)    && throw(ArgumentError("opB must be a unary op!"))
    !is_unary(opC)    && throw(ArgumentError("opC must be a unary op!"))
    !is_unary(opOut)  && throw(ArgumentError("opOut must be a unary op!"))
    descA = CuTensorDescriptor(A; op = opA)
    descB = CuTensorDescriptor(B; op = opB)
    descC = CuTensorDescriptor(C; op = opC)
    # for now, D must be identical to C (and thus, descD must be identical to descC)
    computeType = cutensorComputeType(compute_type)
    T     = sizeof(compute_type) < sizeof(eltype(C)) ? eltype(C) : compute_type
    modeA = collect(Cint, Ainds)
    modeB = collect(Cint, Binds)
    modeC = collect(Cint, Cinds)

    alignmentRequirementA = Ref{UInt32}(C_NULL)
    cutensorGetAlignmentRequirement(handle(), A, descA, alignmentRequirementA)
    alignmentRequirementB = Ref{UInt32}(C_NULL)
    cutensorGetAlignmentRequirement(handle(), B, descB, alignmentRequirementB)
    alignmentRequirementC = Ref{UInt32}(C_NULL)
    cutensorGetAlignmentRequirement(handle(), C, descC, alignmentRequirementC)
    desc = Ref(cutensorContractionDescriptor_t(ntuple(i->0, Val(256))))
    cutensorInitContractionDescriptor(handle(),
                                      desc,
                   descA, modeA, alignmentRequirementA[],
                   descB, modeB, alignmentRequirementB[],
                   descC, modeC, alignmentRequirementC[],
                   descC, modeC, alignmentRequirementC[],
                   computeType)

    find = Ref(cutensorContractionFind_t(ntuple(i->0, Val(64))))
    cutensorInitContractionFind(handle(), find, algo)
    plan = Ref(cutensorContractionPlan_t(ntuple(i->0, Val(640))))
    workspace_size = Ref{UInt64}(C_NULL)
    cutensorContractionGetWorkspace(handle(), desc, find, pref, workspace_size)
    cutensorInitContractionPlan(handle(), plan, desc, find, workspace_size[])
    return plan[]
end

function reduction!(
    alpha::Number, A::CuArray, Ainds::ModeType, opA::cutensorOperator_t,
    beta::Number,  C::CuArray, Cinds::ModeType, opC::cutensorOperator_t,
    opReduce::cutensorOperator_t; stream::CuStream=CuDefaultStream())

    !is_unary(opA)    && throw(ArgumentError("opA must be a unary op!"))
    !is_unary(opC)    && throw(ArgumentError("opC must be a unary op!"))
    !is_binary(opReduce)  && throw(ArgumentError("opReduce must be a binary op!"))
    descA = CuTensorDescriptor(A; op = opA)
    descC = CuTensorDescriptor(C; op = opC)
    # for now, D must be identical to C (and thus, descD must be identical to descC)
    T = eltype(C)
    typeCompute = cutensorComputeType(T)
    modeA = collect(Cint, Ainds)
    modeC = collect(Cint, Cinds)

    @workspace fallback=1<<13 size=@argout(
            cutensorReductionGetWorkspace(handle(),
                A, descA, modeA,
                C, descC, modeC,
                C, descC, modeC,
                opReduce, typeCompute,
                out(Ref{UInt64}(C_NULL)))
        )[] workspace->begin
            cutensorReduction(handle(),
                T[alpha], A, descA, modeA,
                T[beta],  C, descC, modeC,
                        C, descC, modeC,
                opReduce, typeCompute,
                workspace, sizeof(workspace), stream)
        end

    return C
end

function cutensorComputeType(T::DataType)
    if T == Float32
        return CUTENSOR_R_MIN_32F
    elseif T == ComplexF32
        return CUTENSOR_C_MIN_32F
    elseif T == Float16
        return CUTENSOR_R_MIN_16F
    elseif T == ComplexF16
        return CUTENSOR_C_MIN_16F
    elseif T == Float64
        return CUTENSOR_R_MIN_64F
    elseif T == ComplexF64
        return CUTENSOR_C_MIN_64F
    elseif T == Int8
        return CUTENSOR_R_MIN_8I
    elseif T == Int32
        return CUTENSOR_R_MIN_32I
    elseif T == UInt8
        return CUTENSOR_R_MIN_8U
    elseif T == UInt32
        return CUTENSOR_R_MIN_32U
    else
        throw(ArgumentError("cutensorComputeType equivalent for input type $T does not exist!"))
    end
end


using Dates, Random, UUIDs
using Logging

import Base: convert

export
    Model,
    Network,
    InstantNetwork,
    DynamicNetwork,
    Placeholder,
    Node,
    SFunc,
    Dist,
    Score,
    ValueTyped,
    Variable,
    VariableGraph,
    get_initial_graph,
    get_transition_graph,
    get_children,
    get_initial_children,
    get_transition_children,
    get_parents,
    get_initial_parents,
    get_transition_parents,
    get_sfunc,
    get_node,
    get_variables,
    get_placeholders,
    get_nodes,
    input_type,
    make_initial,
    make_transition,
    output_type

"""
    abstract type ValueTyped{O}

Supertype for typing all Scruff Variables; `O` is the
actual type of the variable
"""
abstract type ValueTyped{O} end

"""
value_type(v::ValueTyped{O}) where {O}

return the actual type (i.e. `O`) of the `ValueTyped`
"""
value_type(v::ValueTyped{O}) where {O} = O

"""
    abstract type SFunc{I<:Tuple, O}

A Stochastic Function type with input variables defined by `I` and output type `O`.
This is an abstract representation for a collection of operators with the same
input and output types.

All sfuncs have the following operators defined:
- `compute_lambda`
- `compute_bel`
- `send_pi`
- `outgoing_pis`
- `outgoing_lambdas`

`SFunc` _also_ has both the operators `cpdf` and `logcpdf` defined in terms of the other.
All sfuncs should implement one or the other of these operators.
"""
abstract type SFunc{I<:Tuple, O} end

Base.print_without_params(x::Type{<:SFunc}) = false

"""
    Dist{T} = SFunc{Tuple{}, T}

# Additional supported operators
- `make_factors`
- `send_lambda`
"""
Dist{T} = SFunc{Tuple{}, T}

"""
    Score{I} = SFunc{Tuple{I}, Nothing}

`Score` supports two (2) operators:  `get_score` and `get_log_score`.  `get_log_score`
is defined, by default, using `get_score.`  Every subtype of `Score` _must_ implement
`get_score`.
"""
Score{I} = SFunc{Tuple{I}, Nothing}

"""
    input_type(::Type{<:SFunc{I,O}}) where {I,O}
    input_type(::SFunc{I,O}) where {I,O}

Return the input type (i.e. `I`) of the `SFunc`
"""
input_type(::Type{<:SFunc{I,O}}) where {I,O} = I
input_type(::SFunc{I,O}) where {I,O} = I
"""
    output_type(::Type{<:SFunc{I,O}}) where {I,O}
    output_type(::SFunc{I,O}) where {I,O}

Return the output type (i.e. `O`) of the `SFunc`
"""
output_type(::Type{<:SFunc{I,O}}) where {I,O} = O
output_type(::SFunc{I,O}) where {I,O} = O

"""
    Placeholder{O} <: ValueTyped{O}

A type for typing Scruff variables that do not reference models
"""
struct Placeholder{O} <: ValueTyped{O}
    name::Symbol
end

"""
    Model{I, J, O} <: ValueTyped{O}

Supertype for all Scruff models.
The model represents a variable that varies over time and has output type `O`.
The model may return an initial sfunc with input type `I` using `make_initial`,
which takes the current time as argument,
and a transition sfunc with input type `J` using `make_transition`,
which takes both the parent times (a tuple of times of the same length as `J`)
and the current time as arguments.
These two functions need to be defined for every sfunc.

# Type parameters
- `I`: the input type to the `SFunc` returned by the model's `make_initial` function
- `J`: the input type to the `SFunc` used during the `make_trasition` function call
- `O`: the actual type of the variables represented by this model
"""
abstract type Model{I, J, O} <: ValueTyped{O} end
# is_fixed(m) = false

function make_initial(m,t) end

function make_transition(m, parenttimes, time) end

(model::Model)(symbol::Symbol) = instantiate(model, symbol)

"""
    instantiate(model::Model, name::Symbol)

Create a new `Variable` with the given name for the given model.  Every
`Model` instance is also a function that, given a `Symbol`, will call
this method.
"""
function instantiate(model::Model, name::Symbol)
    return Variable(name, model)
end

"""
    mutable struct Variable{I,J,O} <: ValueTyped{O}

A Variable describes the (time-series) of some set of random values.
It must be named, described by a model, and references to the model inputs must be defined.

For the type variables, see [`Model`](@ref)
"""
mutable struct Variable{I,J,O} <: ValueTyped{O}
    name::Symbol
    model::Model{I,J,O}
end

"""`Node{O} = Union{Placeholder{O}, Variable{I,J,O} where {I,J}}`"""
Node{O} = Union{Placeholder{O}, Variable{I,J,O} where {I,J}}

"""
    output_type(::Variable{I,J,O})
    output_type(::Placeholder{O})

returns the output type `O` of the given parameter
"""
function output_type(::Variable{I,J,O}) where {I,J,O}
    return O
end

function output_type(::Placeholder{O}) where O
    return O
end

get_name(node::Node) = node.name

"""VariableGraph = Dict{Node, Vector{Node}}"""
VariableGraph = Dict{Node, Vector{Node}}

"""
    abstract type Network{I,J,O}

Collects variables, with placeholders for inputs, and defined outputs,
along with up to two directed graphs for initial and transition.
Must implement `get_nodes`, `get_outputs`,
`get_initial_graph`, `get_transition_graph` (which returns a `VariableGraph`)

For the type parameters, see the underlying [`Model`](@ref) class for the mapped
[`Variable`](@ref)s.
"""
abstract type Network{I,J,O} end

"""
    get_initial_parents(n::Network, node::Node)::Vector{Node}

Returns a list of the initial parent nodes in the network from the network's
initial graph (i.e. the return graph from get_initial_graph(network)).
"""
function get_initial_parents(n::Network, node::Node)::Vector{Node}
    g = get_initial_graph(n)
    get(g, node, Node[])
end

"""
    get_transition_parents(n::Network, node::Node)::Vector{Node}

Returns a list of transitioned parent nodes in the network from the network's
transition graph (i.e. the return graph from get_transition_graph(network))
"""
function get_transition_parents(n::Network, node::Node)::Vector{Node}
    g = get_transition_graph(n)
    get(g, node, Node[])
end

"""
    get_node(n::Network, name::Symbol)::Union{Node, Nothing}

Returns the node with the given name, or `nothing`.
"""
function get_node(n::Network, name::Symbol)::Union{Node, Nothing}
    vs = get_nodes(n)
    for v in vs
        if v.name == name
            return v
        end
    end
    return nothing
end

"""
    get_initial_children(n::Network, var::Node)::Vector{Node}

Returns a list of the initial child nodes in the network from the network's
initial graph (i.e. the return graph from get_initial_graph(network)).
"""
function get_initial_children(n::Network, var::Node)::Vector{Node}
    [c for c in get_nodes(n) if var in get_initial_parents(n, c)]
end

"""
    get_transition_children(n::Network, var::Node)::Vector{Node}

Returns a list of transitioned child nodes in the network from the network's
transition graph (i.e. the return graph from get_transition_graph(network))
"""
function get_transition_children(n::Network, var::Node)::Vector{Node}
    [c for c in get_nodes(n) if var in get_transition_parents(n, c)]
end

"""
    complete_graph!(variables::Vector{<:Variable}, graph::VariableGraph)

Populate the `graph` with `v=>Node[]` for all `variables` not in the graph.
"""
function complete_graph!(variables::Vector{<:Variable}, graph::VariableGraph)
    ks = keys(graph)
    for v in variables
        if !(v in ks)
            graph[v] = Node[]
        end
    end
    return graph
end

"""
    struct InstantNetwork{I,J,O} <: Network{I,Nothing,O}

A network that only supports initial models.  The `Nothing` in the supertype
shows that there is no transition to another time.

See [`Network`](@ref) for the type parameters
"""
struct InstantNetwork{I,O} <: Network{I,Nothing,O}
    variables::Vector{<:Variable}
    placeholders::Vector{<:Placeholder}
    outputs::Vector{<:Variable} # outputs is a subset of variables
    parents::VariableGraph
    function InstantNetwork(variables::Vector{<:Variable}, parents::VariableGraph,
                placeholders::Vector{<:Placeholder} = Placeholder[], outputs::Vector{<:Variable} = Variable[])
        I = Tuple{[value_type(p) for p in placeholders]...}
        O = Tuple{[output_type(v) for v in outputs]...}
        return new{I,O}(variables, placeholders, outputs, complete_graph!(variables, parents))
    end
end

get_variables(n::InstantNetwork) = n.variables
get_placeholders(n::InstantNetwork) = n.placeholders
get_nodes(n::InstantNetwork) = union(Set(n.variables), Set(n.placeholders))
get_initial_placeholders(n::InstantNetwork) = n.placeholders
get_transition_placeholders(::InstantNetwork) = error("InstantNetwork does not have a transition graph")
get_outputs(n::InstantNetwork) = n.outputs
get_graph(n::InstantNetwork) = n.parents
get_initial_graph(n::InstantNetwork) = n.parents
get_transition_graph(::InstantNetwork) = error("InstantNetwork does not have a transition graph")
get_parents(n::InstantNetwork, v) = get_initial_parents(n,v)
get_children(n::InstantNetwork, v) = get_initial_children(n,v)

"""
    DynamicNetwork{I,J,O} <: Network{I,J,O}

A network that can transition over time.
See [`Network`](@ref) for the type parameters.
"""
struct DynamicNetwork{I,J,O} <: Network{I,J,O}
    variables::Vector{<:Variable}
    initial_placeholders::Vector{<:Placeholder}
    transition_placeholders::Vector{<:Placeholder}
    outputs::Vector{<:Variable} # outputs is a subset of variables
    initial_parents::VariableGraph
    transition_parents::VariableGraph
    function DynamicNetwork(variables::Vector{<:Variable},
                initial_parents::VariableGraph, transition_parents::VariableGraph,
                init_placeholders::Vector{<:Placeholder} = Placeholder[],
                trans_placeholders::Vector{<:Placeholder} = Placeholder[],
                outputs::Vector{<:Variable} = Variable[])
        I = Tuple{[value_type(p) for p in init_placeholders]...}
        J = Tuple{[value_type(p) for p in trans_placeholders]...}
        O = Tuple{[output_type(v) for v in outputs]...}
        return new{I,J,O}(variables, init_placeholders, trans_placeholders, outputs,
            complete_graph!(variables, initial_parents), complete_graph!(variables, transition_parents))
    end
end

get_variables(n::DynamicNetwork) = n.variables
get_initial_placeholders(n::DynamicNetwork) = n.initial_placeholders
get_transition_placeholders(n::DynamicNetwork) = n.transition_placeholders
get_nodes(n::DynamicNetwork) =
    union(Set(n.variables), Set(n.initial_placeholders), Set(n.transition_placeholders))
get_outputs(n::DynamicNetwork) = n.outputs
get_initial_graph(n::DynamicNetwork) = n.initial_parents
get_transition_graph(n::DynamicNetwork) = n.transition_parents

function assign_times(parents::Dict{Variable, Vector{Node}},
                      timed_varplaces::Dict{Node,Time})::Dict{Node,Time}
    # TODO be smarter. For now this is exhaustive at least, but exponentially slower than needed.
    # Unclear what the exact assumptions on model constraints should be
    # Easy improvement - select immediately identifiable vars with preference (i.e. FixedModel)

    untimed_vars = [var for var in parents.keys() if ~(var in keys(timed_varplaces))]
    if length(untimed_vars)==0
        return timed_varplaces
    end

    for untimed_var in untimed_vars
        new_time = assign_var(parents, timed_varplaces, untimed_var)
        if new_time
            sub_timed_varplaces = copy(timed_varplaces)
            sub_timed_varplaces[untimed_var] = new_time
            recursed = assign_times(parents, sub_timed_varplaces)
        else # Can't assign
            return false
        end

        if recursed
            return recursed
        end # Else try assigning another variable
    end

    return false
end


mutable struct Prop3DAcoVTIDenQ_DEO2_FDTD
    p::Ptr{Cvoid}
end

function Prop3DAcoVTIDenQ_DEO2_FDTD(;
        nz=0,
        ny=0,
        nx=0,
        nsponge=60,
        nbz=512,
        nby=8,
        nbx=8,
        dz=5.0,
        dy=8.0,
        dx=5.0,
        dt=0.001,
        nthreads=Sys.CPU_THREADS,
        freesurface=true,
        freqQ=5.0,
        qMin=0.1,
        qInterior=100.0)
    nz,ny,nx,nsponge,nbz,nby,nbx,nthreads = map(x->round(Int,x), (nz,ny,nx,nsponge,nbz,nby,nbx,nthreads))
    dz,dy,dx,dt = map(x->Float32(x), (dz,dy,dx,dt))

    @assert nx > 0
    @assert ny > 0
    @assert nz > 0
    @assert nsponge > 0
    @assert nthreads > 0
    @assert nbx > 0
    @assert nby > 0
    @assert nbz > 0

    fs = freesurface ? 1 : 0

    p = ccall((:Prop3DAcoVTIDenQ_DEO2_FDTD_alloc, libprop3DAcoVTIDenQ_DEO2_FDTD),
        Ptr{Cvoid},
        (Clong, Clong,    Clong, Clong, Clong, Clong,   Cfloat, Cfloat, Cfloat, Cfloat, Clong, Clong, Clong),
         fs,    nthreads, nx,    ny,    nz,    nsponge, dx,     dy,     dz,     dt,     nbx,   nby,   nbz)

    ccall((:Prop3DAcoVTIDenQ_DEO2_FDTD_SetupDtOmegaInvQ, libprop3DAcoVTIDenQ_DEO2_FDTD),
        Cvoid,
        (Ptr{Cvoid}, Cfloat, Cfloat, Cfloat),
         p,          freqQ,  qMin,   qInterior)

    Prop3DAcoVTIDenQ_DEO2_FDTD(p)
end

free(prop::Prop3DAcoVTIDenQ_DEO2_FDTD) = ccall((:Prop3DAcoVTIDenQ_DEO2_FDTD_free, libprop3DAcoVTIDenQ_DEO2_FDTD), Cvoid, (Ptr{Cvoid},), prop.p)

function size(prop::Prop3DAcoVTIDenQ_DEO2_FDTD)
    nz = ccall((:Prop3DAcoVTIDenQ_DEO2_FDTD_getNz, libprop3DAcoVTIDenQ_DEO2_FDTD), (Clong), (Ptr{Cvoid},), prop.p)
    ny = ccall((:Prop3DAcoVTIDenQ_DEO2_FDTD_getNy, libprop3DAcoVTIDenQ_DEO2_FDTD), (Clong), (Ptr{Cvoid},), prop.p)
    nx = ccall((:Prop3DAcoVTIDenQ_DEO2_FDTD_getNx, libprop3DAcoVTIDenQ_DEO2_FDTD), (Clong), (Ptr{Cvoid},), prop.p)
    (nz,ny,nx)
end

for _f in (:V, :Eps, :Eta, :B, :F, :PSpace, :MSpace, :PCur, :POld, :MCur, :MOld, :DtOmegaInvQ)
    symf = "Prop3DAcoVTIDenQ_DEO2_FDTD_get" * string(_f)
    @eval $(_f)(prop::Prop3DAcoVTIDenQ_DEO2_FDTD) = unsafe_wrap(Array, ccall(($symf, libprop3DAcoVTIDenQ_DEO2_FDTD), Ptr{Float32}, (Ptr{Cvoid},), prop.p), size(prop), own=false)
end

propagateforward_linear!(prop::Prop3DAcoVTIDenQ_DEO2_FDTD) = ccall((:Prop3DAcoVTIDenQ_DEO2_FDTD_TimeStepLinear, libprop3DAcoVTIDenQ_DEO2_FDTD), Cvoid, (Ptr{Cvoid},), prop.p)
propagateforward!(prop::Prop3DAcoVTIDenQ_DEO2_FDTD) = ccall((:Prop3DAcoVTIDenQ_DEO2_FDTD_TimeStep, libprop3DAcoVTIDenQ_DEO2_FDTD), Cvoid, (Ptr{Cvoid},), prop.p)
propagateadjoint!(prop::Prop3DAcoVTIDenQ_DEO2_FDTD) = propagateforward_linear!(prop) # self-adjoint

scale_spatial_derivatives!(prop::Prop3DAcoVTIDenQ_DEO2_FDTD) =
    ccall((:Prop3DAcoVTIDenQ_DEO2_FDTD_ScaleSpatialDerivatives, libprop3DAcoVTIDenQ_DEO2_FDTD), Cvoid, (Ptr{Cvoid},), prop.p)

abstract type Prop3DAcoVTIDenQ_DEO2_FDTD_Model end
struct Prop3DAcoVTIDenQ_DEO2_FDTD_Model_VEA <: Prop3DAcoVTIDenQ_DEO2_FDTD_Model end
struct Prop3DAcoVTIDenQ_DEO2_FDTD_Model_V <: Prop3DAcoVTIDenQ_DEO2_FDTD_Model end

# v,ϵ and η in model-space
function forwardBornInjection!(prop::Prop3DAcoVTIDenQ_DEO2_FDTD, modeltype::Prop3DAcoVTIDenQ_DEO2_FDTD_Model_VEA, dmodel, wavefields)
    ccall((:Prop3DAcoVTIDenQ_DEO2_FDTD_ForwardBornInjection_VEA, libprop3DAcoVTIDenQ_DEO2_FDTD), Cvoid,
        (Ptr{Cvoid}, Ptr{Cfloat}, Ptr{Cfloat}, Ptr{Cfloat}, Ptr{Cfloat},        Ptr{Cfloat},        Ptr{Cfloat},          Ptr{Cfloat}),
         prop.p,     dmodel["v"], dmodel["ϵ"], dmodel["η"], wavefields["pold"], wavefields["mold"], wavefields["pspace"], wavefields["mspace"])
end

function adjointBornAccumulation!(prop::Prop3DAcoVTIDenQ_DEO2_FDTD, modeltype::Prop3DAcoVTIDenQ_DEO2_FDTD_Model_VEA,
        imagingcondition::ImagingConditionStandard, dmodel, wavefields)
    ccall((:Prop3DAcoVTIDenQ_DEO2_FDTD_AdjointBornAccumulation_VEA, libprop3DAcoVTIDenQ_DEO2_FDTD), Cvoid,
        (Ptr{Cvoid}, Ptr{Cfloat}, Ptr{Cfloat}, Ptr{Cfloat}, Ptr{Cfloat},        Ptr{Cfloat},        Ptr{Cfloat},          Ptr{Cfloat}),
         prop.p,     dmodel["v"], dmodel["ϵ"], dmodel["η"], wavefields["pold"], wavefields["mold"], wavefields["pspace"], wavefields["mspace"])
end

# v in model-space
function forwardBornInjection!(prop::Prop3DAcoVTIDenQ_DEO2_FDTD, modeltype::Prop3DAcoVTIDenQ_DEO2_FDTD_Model_V, dmodel, wavefields)
    ccall((:Prop3DAcoVTIDenQ_DEO2_FDTD_ForwardBornInjection_V, libprop3DAcoVTIDenQ_DEO2_FDTD), Cvoid,
        (Ptr{Cvoid}, Ptr{Cfloat}, Ptr{Cfloat},          Ptr{Cfloat}),
         prop.p,     dmodel["v"], wavefields["pspace"], wavefields["mspace"])
end

function adjointBornAccumulation!(prop::Prop3DAcoVTIDenQ_DEO2_FDTD, modeltype::Prop3DAcoVTIDenQ_DEO2_FDTD_Model_V,
        imagingcondition::ImagingConditionStandard, dmodel, wavefields)
    ccall((:Prop3DAcoVTIDenQ_DEO2_FDTD_AdjointBornAccumulation_V, libprop3DAcoVTIDenQ_DEO2_FDTD), Cvoid,
        (Ptr{Cvoid}, Ptr{Cfloat}, Ptr{Cfloat},          Ptr{Cfloat}),
         prop.p,     dmodel["v"], wavefields["pspace"], wavefields["mspace"])
end

function adjointBornAccumulation!(prop::Prop3DAcoVTIDenQ_DEO2_FDTD, modeltype::Prop3DAcoVTIDenQ_DEO2_FDTD_Model_V,
        imagingcondition::ImagingConditionWaveFieldSeparationFWI, dmodel, wavefields)
    ccall((:Prop3DAcoVTIDenQ_DEO2_FDTD_AdjointBornAccumulation_wavefieldsep_V, libprop3DAcoVTIDenQ_DEO2_FDTD), Cvoid,
        (Ptr{Cvoid}, Ptr{Cfloat}, Ptr{Cfloat},          Ptr{Cfloat},          Clong),
         prop.p,     dmodel["v"], wavefields["pspace"], wavefields["mspace"], 1)
end

function adjointBornAccumulation!(prop::Prop3DAcoVTIDenQ_DEO2_FDTD, modeltype::Prop3DAcoVTIDenQ_DEO2_FDTD_Model_V,
        imagingcondition::ImagingConditionWaveFieldSeparationRTM, dmodel, wavefields)
    ccall((:Prop3DAcoVTIDenQ_DEO2_FDTD_AdjointBornAccumulation_wavefieldsep_V, libprop3DAcoVTIDenQ_DEO2_FDTD), Cvoid,
        (Ptr{Cvoid}, Ptr{Cfloat}, Ptr{Cfloat},          Ptr{Cfloat},          Clong),
         prop.p,     dmodel["v"], wavefields["pspace"], wavefields["mspace"], 0)
end

function show(io::IO, prop::Prop3DAcoVTIDenQ_DEO2_FDTD)
    nx = ccall((:Prop3DAcoVTIDenQ_DEO2_FDTD_getNx, libprop3DAcoVTIDenQ_DEO2_FDTD), (Clong), (Ptr{Cvoid},), prop.p)
    ny = ccall((:Prop3DAcoVTIDenQ_DEO2_FDTD_getNy, libprop3DAcoVTIDenQ_DEO2_FDTD), (Clong), (Ptr{Cvoid},), prop.p)
    nz = ccall((:Prop3DAcoVTIDenQ_DEO2_FDTD_getNz, libprop3DAcoVTIDenQ_DEO2_FDTD), (Clong), (Ptr{Cvoid},), prop.p)
    write(io, "Prop3DAcoVTIDenQ_DEO2_FDTD -- nx,ny,nz; $nx,$ny,$nz")
end

function tilesize(::Type{Prop3DAcoVTIDenQ_DEO2_FDTD}; nz, ny, nx, rngy=4:2:64, rngx=4:2:64, N=10, nthreads=Sys.CPU_THREADS-1)
    speeds = Matrix{Float64}(undef, length(rngy), length(rngx))
    for (iby,by) in enumerate(rngy), (ibx,bx) in enumerate(rngx)
        @info "by=$by, bx=$bx"
        prop = Prop3DAcoVTIDenQ_DEO2_FDTD(nz=nz, ny=ny, nx=nx, nbz=nz, nby=by, nbx=bx, nthreads=nthreads)

        v = WaveFD.V(prop)
        ϵ = WaveFD.Eps(prop)
        η = WaveFD.Eta(prop)
        b = WaveFD.B(prop)
        pcur = WaveFD.PCur(prop)
        pold = WaveFD.POld(prop)
        mcur = WaveFD.MCur(prop)
        mold = WaveFD.MOld(prop)

        rand!(pcur)
        rand!(pold)
        rand!(mcur)
        rand!(mold)

        t = 0.0
        set_zero_subnormals(true)
        for i = 1:N
            t += @elapsed propagateforward!(prop)
        end
        set_zero_subnormals(false)
        speeds[iby,ibx] = N*prod(size(prop))/1000/1000/t
        free(prop)
    end
    rngx,rngy,speeds
end


export gate_count

function gate_count(spec)
    ci, _ = code_yao(Semantic.main, spec; optimize = true)
    return gate_count(ci)
end

function gate_count(ci::CodeInfo)
    count = Dict{Symbol,Any}()
    for stmt in ci.code
        is_quantum_statement(stmt) || continue
        qt = quantum_stmt_type(stmt)

        if qt === :gate || qt === :ctrl
            gc = get!(count, qt, IdDict{Any,Int}())
            if stmt.head === :invoke
                gate, gt = obtain_gate_stmt(stmt.args[3], ci)
            elseif stmt.head === :call
                gate, gt = obtain_gate_stmt(stmt.args[2], ci)
            end

            if gate isa IntrinsicRoutine && isempty(gate.variables)
                gc[gate] = get(gc, gate, 0) + 1
            else
                gc[parent(gate)] = get(gc, parent(gate), 0) + 1
            end
        else
            count[qt] = get(count, qt, 0) + 1
        end
    end
    return count
end


include("../src/NMRSpectraSimulator.jl")
import .NMRSpectraSimulator

using LinearAlgebra
using FFTW

import PlotlyJS
using Plots; plotly()

import OffsetArrays
import Interpolations

import BSON
import JSON

import Random
Random.seed!(25)

# simulation compounds. # paths.
plot_title = "resonance groupings"

# name of molecules to simulate.
# Limit to only one compound in this tutorial since we just want to visualize the different resonance groups for one compound.
#molecule_names = ["L-Histidine";]

tag = "L-Isoleucine"
molecule_names = [tag;]

# get mapping from molecule names to their spin system info json files.
H_params_path = "/home/roy/Documents/repo/NMRData/input/coupling_info"
dict_compound_to_filename = JSON.parsefile("/home/roy/Documents/repo/NMRData/input/compound_mapping/select_compounds.json")

# where the bson file is located.
root_folder = "/home/roy/MEGAsync/outputs/NMR/experiments/BMRB-500-0.5mM"
project_path = joinpath(root_folder, tag)
load_path = joinpath(project_path, "experiment.bson")

# where to save the resultant plot. Warning: This script will create the following path if it doesn't it exist.
save_folder = joinpath(project_path, "plots")

# spin-Hamiltonian-related.
tol_coherence = 1e-2 # resonances are pairs of eigenvalues of the Hamiltonian that have quantum numbers that differ by -1. This is the tolerance away from -1 that is allowed.
α_relative_threshold = 0.05 # resonances with relative amplitude less than this factor compared to the maximum resonance in the spin group will be removed. Set to 0.0 to see every single resonance component.
Δc_partition_radius = 0.17 # determines how many resonances get grouped together. Larger number means less groups and thus more resonances per group.
SH_config_path = "/home/roy/Documents/repo/NMRData/input/SH_configs/select_compounds_SH_configs.json"

# surrogate-related.
# default values. The alternative is to load from a config file.
Δr_default = 1.0 # the samples used to build the surrogate is taken every `Δr` radian on the frequency axis. Decrease for improved accuracy at the expense of computation resources.
Δκ_λ_default = 0.05 # the samples used to build thes urrogate for κ_λ are taken at this sampling spacing. Decrease for improved accuracy at the expense of computation resources.
Δcs_max_scalar_default = 0.2 # In units of ppm. interpolation border that is added to the lowest and highest resonance frequency component of the mixture being simulated.
κ_λ_lb_default = 0.5 # interpolation lower limit for κ_λ.
κ_λ_ub_default = 2.5 # interpolation upper limit for κ_λ.
λ0 = 3.4
surrogate_config_path = "/home/roy/Documents/repo/NMRData/input/surrogate_configs/select_compounds_surrogate_configs.json"

dummy_SSFID = NMRSpectraSimulator.SpinSysParamsType1(0.0) # level 2 model.

##### end user inputs.

### load block.
dic = BSON.load(load_path)
fs = dic[:fs]
SW = dic[:SW]
ν_0ppm = dic[:ν_0ppm]
λ_0ppm = dic[:λ_0ppm]

hz2ppmfunc = uu->(uu - ν_0ppm)*SW/fs
ppm2hzfunc = pp->(ν_0ppm + pp*fs/SW)

# get a surrogate where K_{n,i} is encouraged to be no larger than `early_exit_part_size`.
#println("Timing: setupmixtureproxies()")
#@time mixture_params = NMRSpectraSimulator.setupmixtureproxies(molecule_names,
mixture_params = NMRSpectraSimulator.setupmixtureSH(molecule_names,
    H_params_path, dict_compound_to_filename, fs, SW,
    ν_0ppm;
    config_path = SH_config_path,
    tol_coherence = tol_coherence,
    α_relative_threshold = α_relative_threshold,
    Δc_partition_radius = Δc_partition_radius)
As = mixture_params

# We only work with a single compound in this tutorial. Assign a new object for this compound to reduce clutter.
A = As[1];

## frequency locations. For plotting.
ΩS_ppm = NMRSpectraSimulator.getPsnospininfo(mixture_params, hz2ppmfunc)
ΩS_ppm_sorted = sort(NMRSpectraSimulator.combinevectors(ΩS_ppm))

u_offset = 0.5 # in units ppm.
u_min = ppm2hzfunc(ΩS_ppm_sorted[1] - u_offset)
u_max = ppm2hzfunc(ΩS_ppm_sorted[end] + u_offset)

# This is the frequency range that we shall work with.
P = LinRange(hz2ppmfunc(u_min), hz2ppmfunc(u_max), 50000)
U = ppm2hzfunc.(P)
U_rad = U .* (2*π)

# fit the surrogate.
Bs = NMRSpectraSimulator.fitproxies(As, dummy_SSFID, λ0;
    names = molecule_names,
    config_path = surrogate_config_path,
    Δcs_max_scalar_default = Δcs_max_scalar_default,
    κ_λ_lb_default = κ_λ_lb_default,
    κ_λ_ub_default = κ_λ_ub_default,
    u_min = u_min,
    u_max = u_max,
    Δr_default = Δr_default,
    Δκ_λ_default = Δκ_λ_default)
B = Bs[1] # looking at the first (and only) molecule in this example script.

# create the functions for each resonance group.
qs = collect( collect( ωω->B.qs[i][k](ωω-B.ss_params.d[i], B.ss_params.κs_λ[i]) for k = 1:length(B.qs[i]) ) for i = 1:length(B.qs) )
q_singlets = ωω->NMRSpectraSimulator.evalsinglets(ωω, B.d_singlets, A.αs_singlets, A.Ωs_singlets, B.β_singlets, B.λ0, B.κs_λ_singlets)

# create the function for the entire compound.
q = uu->NMRSpectraSimulator.evalitpproxymixture(uu, As[1:1], Bs[1:1])

# evaluate at the plotting positions.
q_U = q.(U_rad)

qs_U = collect( collect( qs[i][k].(U_rad) for k = 1:length(qs[i]) ) for i = 1:length(qs) )
q_singlets_U = q_singlets.(U_rad)

# sanity check.
q_check_U = q_singlets_U
if !isempty(qs) # some compounds only have singlets.
    q_check_U += sum( sum( qs[i][k].(U_rad) for k = 1:length(qs[i]) ) for i = 1:length(qs) )
end
discrepancy = norm(q_check_U- q_U)
println("sanity check. This should be numerically zero: ", discrepancy)

"""
save resonance groupings of a compound.
Choices for `f`` are: `real()`, `imag()`, or `abs()`. These corresponds to real part, imaginary part, and magnitude spectrum, respectively.
"""
function plotgroups(title_string::String,
    P,
    U,
    q,
    qs,
    q_singlets,
    f::Function,
    return_val_type::T;
    canvas_size::Tuple{Int,Int} = (1000, 400)) where T

    U_rad = U .* (2*π)

    q_U = q.(U_rad)
    plot_obj = Plots.plot( P,
        f.(q_U),
        title = title_string,
        label = "sum of sub-models",
        seriestype = :line,
        ticks = :native,
        xlims = (P[1],P[end]),
        hover = P,
        linewidth = 4,
        xflip = true,
        size = canvas_size)

    qs_U = Vector{Vector{Vector{Complex{T}}}}(undef, length(qs))
    for i = 1:length(qs)

        qs_U[i] = Vector{Vector{Complex{T}}}(undef, length(qs[i]))
        for k = 1:length(qs[i])

            qs_U[i][k] = qs[i][k].(U_rad)

            Plots.plot!(plot_obj, P, f.(qs_U[i][k]), label = "sub-group ($(i),$(k))",
                seriestype = :line,
                linestyle = :dot,
                xflip = true,
                linewidth = 4)
        end
    end

    q_singlets_U = q_singlets.(U_rad)
    Plots.plot!(plot_obj, P, f.(q_singlets_U), label = "group of all singlets",
        seriestype = :line,
        linestyle = :dot,
        xflip = true,
        linewidth = 4)
    #
    #Plots.plot!(plot_obj, P, f.(q_U),
    #markershape = :circle,
    #seriestype = :scatter,
    #xflip = true)

    return plot_obj, q_U, qs_U, q_singlets_U
end

# reduce the plotting positions for low signal regions. Otherwise the plot store size will be too large, and the time to load the plot will be long.
display_reduction_factor = 100
display_threshold_factor =  α_relative_threshold/10

inds, _ = NMRSpectraSimulator.prunelowsignalentries(q_U, display_threshold_factor, display_reduction_factor)
P_display = P[inds]
U_display = U[inds]

# plot.
canvas_size = (1000, 400)

plots_save_path = joinpath(save_folder, "groups_real.html")
title_string = "$(plot_title), real"
plot_obj, q_U, qs_U, q_singlets_U = plotgroups(title_string, P_display, U_display, q, qs, q_singlets, real, P[1]; canvas_size = canvas_size)
Plots.savefig(plot_obj, plots_save_path)
display(plot_obj)


module Land

using DiffEqOperators
using DocStringExtensions
using OrdinaryDiffEq
using Parameters
using PkgUtility: T_25
using WaterPhysics: TraceGasAir, TraceGasCO₂, diffusive_coefficient

# export public types
export VerticalLayers

# export public functions
export vertical_layers!

# include sub modules
include("CanopyLayers/CanopyLayers.jl"                  )
include("Photosynthesis/Photosynthesis.jl"              )
include("PlantHydraulics/PlantHydraulics.jl"            )
include("StomataModels/StomataModels.jl"                )
include("SoilPlantAirContinuum/SoilPlantAirContinuum.jl")

# include types
include("Land/Types.jl"    )
include("Land/AirLayers.jl")

end # module