module Microclimate

using Unitful, NCDatasets
using Unitful: W, m, °C, hr, mol, K, s, J, Mg, g, kg, L, kPa, Pa

export AbstractEnvironment,
       AbstractMicroclimEnvironment,
       AbstractMicroclimGrid, MicroclimGrid,
       AbstractMicroclimPoint, MicroclimPoint,
       AbstractMicroclimInstant, MicroclimInstant,
       LinearLayerInterpolator, LinearLayerInterpolators

export radiation, snowdepth, airtemperature, relhumidity, windspeed,
       soiltemperature, soilwaterpotential, soilwatercontent

export mean_soiltemperature, mean_soilwaterpotential, mean_soilwatercontent

export weightedmean, layermax

export load_grid, load_point

radiation_type = typeof(0.1W*m^-2)
snowdepth_type = typeof(0.1m)
airtemperature_type = typeof(0.1K)
relhumidity_type = typeof(0.001)
windspeed_type = typeof(0.1m*s^-1)
soiltemperature_type = typeof(0.1K)
soilwaterpotential_type = typeof(0.0kPa)
soilwatercontent_type = typeof(0.001)

abstract type AbstractEnvironment end
abstract type AbstractMicroclimEnvironment{S,I,A,R} <: AbstractEnvironment end
abstract type AbstractMicroclimGrid{S,I,A,R} <: AbstractMicroclimEnvironment{S,I,A,R} end
abstract type AbstractMicroclimPoint{S,I,A,R} <: AbstractMicroclimEnvironment{S,I,A,R} end

struct MicroclimGrid{S,I,A,R,RA,SN,AT,RH,WS,ST,WP,WC} <: AbstractMicroclimGrid{S,I,A,R}
    radiation::RA
    snowdepth::SN
    airtemperature::AT
    relhumidity::RH
    windspeed::WS
    soiltemperature::ST
    soilwaterpotential::WP
    soilwatercontent::WC

    MicroclimGrid{S,I,A,R}(radiation::RA, snowdepth::SN, airtemperature::AT, relhumidity::RH,
                            windspeed::WS, soiltemperature::ST, soilwaterpotential::WP,
                            soilwatercontent::WC) where {S,I,A,R,RA,SN,AT,RH,WS,ST,WP,WC} =
        new{S,I,A,R,RA,SN,AT,RH,WS,ST,WP,WC}(radiation, snowdepth, airtemperature,
            relhumidity, windspeed, soiltemperature, soilwaterpotential, soilwatercontent)

    MicroclimGrid{S,I,A,R,RA,SN,AT,RH,WS,ST,WP,WC}(radiation::RA, snowdepth::SN, airtemperature::AT, relhumidity::RH,
                            windspeed::WS, soiltemperature::ST, soilwaterpotential::WP,
                            soilwatercontent::WC) where {S,I,A,R,RA,SN,AT,RH,WS,ST,WP,WC} =
        new{S,I,A,R,RA,SN,AT,RH,WS,ST,WP,WC}(radiation, snowdepth, airtemperature,
            relhumidity, windspeed, soiltemperature, soilwaterpotential, soilwatercontent)
end

struct MicroclimPoint{S,I,A,R,RA,SN,AT,RH,WS,ST,WP,WC} <: AbstractMicroclimPoint{S,I,A,R}
    radiation::RA
    snowdepth::SN
    airtemperature::AT
    relhumidity::RH
    windspeed::WS
    soiltemperature::ST
    soilwaterpotential::WP
    soilwatercontent::WC

    MicroclimPoint{S,I,A,R}(radiation::RA, snowdepth::SN, airtemperature::AT, relhumidity::RH,
                            windspeed::WS, soiltemperature::ST, soilwaterpotential::WP,
                            soilwatercontent::WC) where {S,I,A,R,RA,SN,AT,RH,WS,ST,WP,WC} =
        new{S,I,A,R,RA,SN,AT,RH,WS,ST,WP,WC}(radiation, snowdepth, airtemperature,
            relhumidity, windspeed, soiltemperature, soilwaterpotential, soilwatercontent)

    MicroclimPoint{S,I,A,R,RA,SN,AT,RH,WS,ST,WP,WC}(radiation::RA, snowdepth::SN, airtemperature::AT, relhumidity::RH,
                            windspeed::WS, soiltemperature::ST, soilwaterpotential::WP,
                            soilwatercontent::WC) where {S,I,A,R,RA,SN,AT,RH,WS,ST,WP,WC} =
        new{S,I,A,R,RA,SN,AT,RH,WS,ST,WP,WC}(radiation, snowdepth, airtemperature,
            relhumidity, windspeed, soiltemperature, soilwaterpotential, soilwatercontent)
end

abstract type AbstractMicroclimInstant end

struct MicroclimInstant{M,I,T} <: AbstractMicroclimInstant
    microclimate::M
    interp::I
    t::T
end
MicroclimInstant(microclim, height::Number, t) = MicroclimInstant(microclim, LinearLayerInterpolators(microclim, height), t)

# From external data without units
MicroclimPoint{S,I,A,R}(radiation::Vector{<:AbstractFloat}, snowdepth::Vector{<:AbstractFloat},
                      airtemperature::Array{<:AbstractFloat}, relhumidity::Array{<:AbstractFloat},
                      windspeed::Array{<:AbstractFloat}, soiltemperature::Array{<:AbstractFloat},
                      soilwaterpotential::Array{<:AbstractFloat}, soilwatercontent::Array{<:AbstractFloat}
                     ) where {S,I,A,R} = begin

    rad = to_radiation.(radiation)
    snow = to_snowdepth.(snowdepth)
    airt = to_airtemperature.(airtemperature)
    relh = to_relhumidity.(relhumidity)
    wind = to_windspeed.(windspeed)
    soilt = to_soiltemperature.(soiltemperature)
    soilwp = to_soilwaterpotential.(soilwaterpotential)
    soilwc = to_soilwatercontent.(soilwatercontent)

    args = (rad, snow, airt, relh, wind, soilt, soilwp, soilwc)
    MicroclimPoint{I,A,R,typeof.(args)...}(args...)
end

# From microclim dataset in tuples and a cartesian index
# Requires conversion to float from Int
MicroclimPoint{S,I,A,R}(radiation::Vector, snowdepth::Vector, airtemperature::Tuple,
                        relhumidity::Tuple, windspeed::Tuple, soiltemperature::Tuple,
                        soilwaterpotential::Tuple, soilwatercontent::Tuple, i::CartesianIndex) where {S,I,A,R} = begin
    rad = to_radiation.(catpoint(radiation, i) .* 0.1)
    snow = to_snowdepth.(catpoint(snowdepth, i) .* 0.001)
    airt = to_airtemperature.(catpoint(airtemperature, i) .* 0.1)
    relh = to_relhumidity.(catpoint(relhumidity, i) .* 0.001)
    wind = to_windspeed.(catpoint(windspeed, i) .* 0.1)
    soilt = to_soiltemperature.(catpoint(soiltemperature, i) .* 0.1)
    soilwp = to_soilwaterpotential.(catpoint(soilwaterpotential, i) .* 0.1)
    soilwc = to_soilwatercontent.(catpoint(soilwatercontent, i) .* 0.001)

    args = (rad, snow, airt, relh, wind, soilt, soilwp, soilwc)
    MicroclimPoint{S,I,A,R,typeof.(args)...}(args...)
end

# From MiriclimGrid and a cartesian index
MicroclimPoint(env::AbstractMicroclimGrid{S,I,A,R}, i::CartesianIndex) where {S,I,A,R} =
    MicroclimPoint{S,I,A,R}(radiation(env), snowdepth(env), airtemperature(env),
                   relhumidity(env), windspeed(env), soiltemperature(env),
                   soilwaterpotential(env), soilwatercontent(env), i)

struct MicroclimateLayers{L<:Tuple}
    layers::L
end

include("interpolation.jl")
include("netcdf.jl")

# unit conversions
to_radiation(x) = x * W*m^-2
to_snowdepth(x) = x * m
to_airtemperature(x) = x * °C |> K
to_relhumidity(x) = x
to_windspeed(x) = x * m*s^-1
to_soiltemperature(x) = x * °C |> K
to_soilwaterpotential(x) = x * kPa
to_soilwatercontent(x) = x

# Get layer specifications from type system
get_shade(layers::AbstractMicroclimEnvironment{S,I,A,R}) where {S,I,A,R} = S
get_increments(layers::AbstractMicroclimEnvironment{S,I,A,R}) where {S,I,A,R} = I
get_nextlayer(layers::AbstractMicroclimEnvironment{S,I,A,R}) where {S,I,A,R} = A
get_range(layers::AbstractMicroclimEnvironment{S,I,A,R}) where {S,I,A,R} = R

@inline catpoint(layers::Tuple{}, i::CartesianIndex) = Float64[]
@inline catpoint(layers::Tuple, i::CartesianIndex) =
    cat([[cat((getindex(y, i, :) for y in l)..., dims=1)...] for l in layers]..., dims=2)
@inline catpoint(layers::Vector, i::CartesianIndex) = cat((getindex(y, i, :) for y in layers)..., dims=1)

# Accessors for single layer data
for name in (:radiation, :snowdepth)
    eval(quote
        $name(env::AbstractEnvironment) = env.$name
        $name(env::AbstractMicroclimPoint, interp, i) = lin_interp($name(env), i)
        $name(env::AbstractMicroclimPoint, i) = lin_interp($name(env.microclimate), i)
        $name(env::AbstractMicroclimInstant) = lin_interp($name(env.microclimate), env.t)
    end)
end

# Accessor functions for range and increment data
for (name, mode) in ((:airtemperature, :range), (:relhumidity, :range), (:windspeed, :range),
                     (:soiltemperature, :increment), (:soilwaterpotential, :increment),
                     (:soilwatercontent, :increment))
    eval(quote
        $name(env::AbstractMicroclimEnvironment) = env.$name
        $name(env::AbstractMicroclimPoint, interp::LinearLayerInterpolators, i) =
            $name(env::AbstractMicroclimPoint, interp.$mode, i)
        $name(env::AbstractMicroclimPoint, interp::LinearLayerInterpolator, i) =
            interp_layer($name(env), interp, i)
        $name(env::AbstractMicroclimPoint, height::Number, i) =
            $name(env, $(Symbol(mode, :_interpolator))(height), i)
        $name(env::AbstractMicroclimInstant) =
            $name(env.microclimate, env.interp, env.t)
        $name(env::AbstractMicroclimInstant, i::Number) =
            $name(env.microclimate, env.interp, i)
        $name(env::AbstractMicroclimInstant, interp::AbstractLayerInterpolator) =
            $name(env.microclimate, interp, env.t)
    end)
end

# Accessors for weighted mean of 8 layer data
for name in (:soiltemperature, :soilwatercontent, :soilwaterpotential)
    eval(quote
        $(Symbol(:mean_, name))(env::AbstractMicroclimPoint, x, i) =
            weightedmean(env, $name(env), x, i)
        $(Symbol(:mean_, name))(env::AbstractMicroclimInstant) =
            weightedmean(env, $name(env.microclimate), env.interp, env.t)
        $(Symbol(:mean_, name))(env::AbstractMicroclimInstant, i::Number) =
            weightedmean(env, $name(env.microclimate), env.interp, i)
        $(Symbol(:mean_, name))(env::AbstractMicroclimInstant, interp::AbstractLayerInterpolator) =
            weightedmean(env, $name(env.microclimate), interp, env.t)
    end)
end

end # module


#LIKELIHOOD SCORING FUNCS
function IPM_likelihood(sources::AbstractVector{<:Tuple{<:AbstractMatrix{<:AbstractFloat},<:Integer}}, observations::AbstractMatrix{<:Integer}, obs_lengths::AbstractVector{<:Integer}, bg_scores::AbstractArray{<:AbstractFloat}, mix::BitMatrix, revcomp::Bool=REVCOMP, returncache::Bool=false, cache::AbstractVector{<:AbstractFloat}=zeros(0), clean::AbstractVector{<:Bool}=Vector(falses(size(observations)[2])))
    source_wmls=[size(source[1])[1] for source in sources]
    O = size(bg_scores)[2]
    source_stops=[obsl-wml+1 for wml in source_wmls, obsl in obs_lengths] #stop scannng th source across the observation as the source reaches the end
    L=maximum(obs_lengths)+1

    obs_src_idxs=mix_pull_idxs(mix) #get vectors of sources emitting in each obs

    revcomp ? (srcs=[cat(source[1],revcomp_pwm(source[1]),dims=3) for source in sources]; motif_expectations = [((MOTIF_EXPECT/2)/obsl) for obsl in obs_lengths]; mat_dim=2) : (srcs=[source[1] for source in sources]; ; motif_expectations = [(MOTIF_EXPECT/obsl) for obsl in obs_lengths]; mat_dim=1) #setup appropriate reverse complemented sources if necessary and set log_motif_expectation-nMica has 0.5 per base for including the reverse complement, 1 otherwise

    lme_vec=zeros(length(sources))

    obs_lhs=Vector{Vector{Float64}}() #setup likelihood vecs for threaded operation-reassembled on return
    nt=Threads.nthreads()
    for t in 1:nt-1
        push!(obs_lhs,zeros(Int(floor(O/nt))))
    end
    push!(obs_lhs, zeros(Int(floor(O/nt)+(O%nt))))

    Threads.@threads for t in 1:nt
        revcomp && (weavevec=zeros(3))
        revcomp ? (score_mat=zeros(maximum(source_stops),2)) : (score_mat=zeros(maximum(source_stops)))
        opt = floor(O/nt) #obs per thread
        score_matrices=Vector{typeof(score_mat)}(undef, length(sources)) #preallocate
        osi_emitting=Vector{Int64}() #preallocate
        lh_vec = zeros(L) #preallocated likelihood vector is one position (0 initialiser) longer than the longest obs

        for i in 1:Int(opt+(t==nt)*(O%nt))
            o=Int(i+(t-1)*opt)
            if clean[o]
                obs_lhs[t][i]=cache[o]
            else
                obsl = obs_lengths[o]
                oidxs=obs_src_idxs[o]
                mixwmls=source_wmls[oidxs]

                obs_cardinality = length(oidxs) #the more sources, the greater the cardinality_penalty
                if obs_cardinality > 0
                    revcomp ? score_sources_ds!(score_mat, score_matrices, view(srcs,oidxs), view(observations,:,o), view(source_stops,oidxs,o)) :
                         score_sources_ss!(score_mat, score_matrices, view(srcs,oidxs), view(observations,:,o), view(source_stops,oidxs,o)) #get scores for this observation

                    lme_vec.=motif_expectations[o]
                    penalty_sum = sum(lme_vec[1:obs_cardinality])
                    penalty_sum > 1. && (penalty_sum=1.)
                    cardinality_penalty=log(1.0-penalty_sum)
                else
                    cardinality_penalty=0.0
                end

                revcomp ? (obs_lhs[t][i]=weave_scores_ds!(weavevec, lh_vec, obsl, view(bg_scores,:,o), score_matrices, oidxs, mixwmls, log(motif_expectations[o]), cardinality_penalty, osi_emitting)) :
                    (obs_lhs[t][i]=weave_scores_ss!(lh_vec, obsl, view(bg_scores,:,o), score_matrices, oidxs, mixwmls, log(motif_expectations[o]), cardinality_penalty, osi_emitting))

                empty!(osi_emitting)

            end
        end
    end

    returncache ? (return lps([lps(obs_lhs[t]) for t in 1:nt]), vcat(obs_lhs...)) : (return lps([lps(obs_lhs[t]) for t in 1:nt]))
end
                @inline function mix_pull_idxs(A::AbstractArray{Bool})
                    n=count(A)
                    S=[Vector{Int64}() for o in 1:size(A,1)]
                    cnt=1
                    for (i,a) in pairs(A)
                        if a
                            push!(S[i[1]],i[2])
                            cnt+=1
                        end
                    end
                    return S
                end

                @inline function revcomp_pwm(pwm::AbstractMatrix{<:AbstractFloat}) #in order to find a motif on the reverse strand, we scan the forward strand with the reverse complement of the pwm, reordered 3' to 5', so that eg. an PWM for an ATG motif would become one for a CAT motif
                    return pwm[end:-1:1,end:-1:1]
                end

                @inline function score_sources_ds!(score_mat, score_matrices::AbstractVector{<:AbstractMatrix{<:AbstractFloat}}, sources, observation::AbstractVector{<:Integer}, source_stops)
                    for (s,source) in enumerate(sources)
                        for t in 1:source_stops[s]
                            for position in 1:size(source,1)
                                score_loc = t+position-1 #score_loc is the position of the obs to be scored by PWM
                                score_mat[t,1] += source[position,observation[score_loc],1] #add the appropriate log PWM value from the source to the score
                                score_mat[t,2] += source[position,observation[score_loc],2] #add the appropriate log PWM value from the source to the score

                            end
                        end
                        score_matrices[s]=score_mat[1:source_stops[s],:] #copy score matrix to vector
                        score_mat.=0. #reset score matrix
                    end
                end

                @inline function score_sources_ss!(score_mat, score_matrices::AbstractVector{<:AbstractArray{<:AbstractFloat}}, sources, observation::AbstractVector{<:Integer}, source_stops)
                    for (s,source) in enumerate(sources)
                        for t in 1:source_stops[s]
                            for position in 1:size(source,1)
                                score_loc = t+position-1 #score_loc is the position of the obs to be scored by PWM
                                score_mat[t] += source[position,observation[score_loc]] #add the appropriate log PWM value from the source to the score

                            end
                        end
                        score_matrices[s]=score_mat[1:source_stops[s]] #copy score matrix to vector
                        score_mat.=0. #reset score matrix
                    end
                end

                @inline function weave_scores_ds!(weavevec, lh_vec, obsl::Integer, bg_scores::SubArray, score_mat::AbstractVector{<:AbstractMatrix{<:AbstractFloat}}, obs_source_indices::AbstractVector{<:Integer}, source_wmls::AbstractVector{<:Integer}, log_motif_expectation::AbstractFloat, cardinality_penalty::AbstractFloat, osi_emitting)
                    for i in 2:obsl+1 #i=1 is ithe lh_vec initializing 0, i=2 is the score of the first background position (ie t=1)
                        t=i-1
                        score = lps(lh_vec[i-1], bg_scores[t], cardinality_penalty)

                        #logic: all observations are scored from t=wml to the end of the obs-therefore check at each position for new sources to add (indexed by vector position to retrieve source wml and score matrix)
                        if length(osi_emitting)<length(obs_source_indices)
                            for n in 1:length(obs_source_indices)
                                if !(n in osi_emitting)
                                    t>= source_wmls[n] && (push!(osi_emitting,n))
                                end
                            end
                        end

                        for n in osi_emitting
                            wml = source_wmls[n]
                            from_score = lh_vec[i-wml+1] #score at the first position of the PWM
                            score_array = score_mat[n] #get the source score matrix
                            score_idx = t - wml + 1 #translate t to score_array idx for emission score
                            f_emit_score = score_array[score_idx,1] #emission score at the last position of the PWM
                            r_emit_score = score_array[score_idx,2]

                            weavevec .= score, lps(from_score, f_emit_score, log_motif_expectation), lps(from_score, r_emit_score, log_motif_expectation)

                            score=logsumexp(weavevec)
                        end
                        lh_vec[i] = score
                    end
                    return lh_vec[obsl+1]
                end

                @inline function weave_scores_ss!(lh_vec, obsl::Integer, bg_scores::SubArray, score_mat::AbstractVector{<:AbstractVector{<:AbstractFloat}}, obs_source_indices::AbstractVector{<:Integer}, source_wmls::AbstractVector{<:Integer}, log_motif_expectation::AbstractFloat, cardinality_penalty::AbstractFloat, osi_emitting)
                    for i in 2:obsl+1 #i=1 is ithe lh_vec initializing 0, i=2 is the score of the first background position (ie t=1)
                        t=i-1
                        score = lps(lh_vec[i-1], bg_scores[t], cardinality_penalty)
                        #logic: all observations are scored from t=wml to the end of the obs-therefore check at each position for new sources to add (indexed by vector position to retrieve source wml and score matrix)
                        if length(osi_emitting)<length(obs_source_indices)
                            for n in 1:length(obs_source_indices)
                                if !(n in osi_emitting)
                                    t>= source_wmls[n] && (push!(osi_emitting,n))
                                end
                            end
                        end

                        for n in osi_emitting
                            wml = source_wmls[n]
                            from_score = lh_vec[i-wml+1] #score at the first position of the PWM
                            score_array = score_mat[n] #get the source score matrix
                            score_idx = t - wml + 1 #translate t to score_array idx for emission score
                            f_emit_score = score_array[score_idx] #emission score at the last position of the PWM

                            score=logaddexp(score, lps(from_score, f_emit_score, log_motif_expectation))
                        end
                        lh_vec[i] = score
                    end
                    return lh_vec[obsl+1]
                end


###################################
## Simulation models ##
###################################

function count_active_agents(agents::Vector{Agent})
    r = 0
    @simd for i = eachindex(agents)
        @inbounds r += agents[i].active
    end
    r
end

"""
`simulation_run` run traffic simulation in two modes.

**Input parameters**
* `mode` : simulation model switch
    * `:base` : simulation with regular agents only
    * `:smart` : simulation with smart agents enabled
* `OSMmap` : OpenStreetMapX MapData object with road network data
* `inAgents` : set of agents created with generate_agents function
* `density_factor` : road length reserved for one vehicle
* `k_routes_dict` : dictionary with multiple shortest paths (values) between vertices (keys)
* `U` : period of weights updates
* `T` : distribution parameter in k-shortest path rerouting
* `k` : number of fastest routes generated in rerouting function
* `debug` : debug messages switch
* `track_avg_speeds` : return average speeds on edges if true
"""
function simulation_run(mode::Symbol,
                        OSMmap::MapData,
                        inAgents::Vector{Agent},
                        density_factor::Float64 = 5.0,
                        k_routes_dict::Dict{Tuple{Int,Int},Array{Vector{Int}}}=
                        Dict{Tuple{Int,Int},Array{Vector{Int}}}(),
                        U::Int = 100,
                        T::Float64 = 1.0,
                        k::Int = 3,
                        debug::Bool = false;
                        track_avg_speeds::Bool = false)
    if !in(mode, [:base,:smart]) error("Wrong mode specified.") end
    Agents = deepcopy(inAgents) #Creating working copy of agents
    Agents[1].start_time = 0.0
    Agents[1].active = true
    max_densities, max_speeds = traffic_constants(OSMmap, density_factor) #Traffic characteristic constants
    densities, speeds = init_traffic_variables(OSMmap, Agents) #Traffic characteristic variables
    update_weights!(speeds, densities, max_densities, max_speeds) #Initial speeds update
    curr_edges = [a.edge for a in Agents] #Initialize edges list
    if track_avg_speeds
        avg_speeds = deepcopy(speeds)
        tick = 1
    end
    #Initialize simulation variables
    simtime = 0.0
    steps = 0
    active = length(Agents)
    #Calculate initial time to nearest junction for all agents
    times_to_edge = next_edge(Agents, speeds, OSMmap.w)
    times_to_start = [a.start_time for a in Agents] #Get agents starting times
    times_to_start[1] = Inf
    event_update = Inf
    #Loop until all agents are deactivated
    while active != 0 || !all(times_to_start.==Inf)
        steps += 1
        #Calculate next event time and update sim clock
        event_edge, IDe = findmin(times_to_edge)
        event_start, IDs = findmin(times_to_start)
        if mode == :smart
            event_update = next_update = (simtime ÷ U + 1) * U - simtime
        end
        min_event, event_type = findmin([event_edge, #1
                                        event_start, #2
                                        event_update]) #3
        simtime += min_event #Increase total simulation time
        times_to_edge .-= min_event #Move agents forward to event time
        times_to_start .-= min_event
        if event_type == 1 #Switching edges
            vAgent = Agents[IDe] #Pick agent for which event is occuring
            #Process agent for which event is occuring
            changed_edges = update_event_agent!(vAgent, simtime, densities, OSMmap.v)
            length(changed_edges) == 1 ? curr_edges[IDe] = (0,0) : curr_edges[IDe] = changed_edges[2]
            #Update speeds and correct events time
            update_weights_and_events!(Agents, curr_edges, times_to_edge, speeds,
                changed_edges, densities, max_densities, max_speeds)
            vAgent.active ? times_to_edge[IDe] = next_edge(vAgent, speeds, OSMmap.w) : times_to_edge[IDe] = Inf
        elseif event_type == 2 #New agent joins traffic
            Agents[IDs].active = true
            times_to_start[IDs] = Inf
            times_to_edge[IDs] = next_edge(Agents[IDs], speeds, OSMmap.w)
            c_edge = Agents[IDs].edge
            haskey(densities, c_edge) ? densities[c_edge] += 1 : densities[c_edge] = 1
        elseif event_type == 3 #Velocities update
            if debug
                update = Int(simtime ÷ U)
                smartAgents = sum([a.smart*a.active for a in Agents])
                print("Update: $(update) | Smart agents: $(smartAgents) ")
            end
            #Process rerouting for smart agents
            for a in Agents
                a.smart && a.active &&
                k_shortest_path_rerouting!(OSMmap, k_routes_dict, a, speeds, max_speeds, k, T, U)
            end
            debug && println("Finished")
        end
        #Update average speeds every 30 secs
        if track_avg_speeds && simtime > tick*30
            #Iterative mean
            tick +=1
            avg_speeds = ((tick-1)/tick)*avg_speeds+(1/tick)*speeds
        end
        active = count_active_agents(Agents)
    end
    times = getfield.(Agents,:travel_time)
    if track_avg_speeds
        output_tuple = (Steps = steps,
                    Simtime = simtime,
                    TravelTimes = times,
                    AvgSpeeds = avg_speeds)
    else
        output_tuple = (Steps = steps,
                    Simtime = simtime,
                    TravelTimes = times)
    end
    return output_tuple
end

"""
`gather_statistics` return various measures describing TMS model quality

**Input parameters**
* `smart_ind` : vector with smart agent indicators
* `times_base` : agents travelling time in base scenario
* `times_smart` : agents travelling time in smart scenario
"""
function gather_statistics(smart_ind::BitArray{1},
                            times_base::Vector{Float64},
                            times_smart::Vector{Float64})
    overall_time = (sum(times_base) - sum(times_smart))/sum(times_base)
    smart_time = (sum(times_base[smart_ind]) - sum(times_smart[smart_ind]))/sum(times_base[smart_ind])
    other_time = (sum(times_base[.!smart_ind]) - sum(times_smart[.!smart_ind]))/sum(times_base[.!smart_ind])
    statistics_tuple = (
        overall_time = round(overall_time, digits=3),
        smart_time = round(smart_time, digits=3),
        other_time = round(other_time, digits=3),
        avg_base = round(StatsBase.mean(times_base), digits=2),
        avg_overall_V2I = round(StatsBase.mean(times_smart), digits=2),
        avg_smart_V2I = round(StatsBase.mean(times_smart[smart_ind]), digits=2),
        avg_regular_V2I = round(StatsBase.mean(times_smart[.!smart_ind]), digits=2)
                        )
    return statistics_tuple
end

"""
`run_parameter_analysis` run base and V2I simulation scenario to compare agents performance, returns comparison statistics

**Input parameters**
* `GridElement` : element from parameters grid to be used in simulation
* `ParamGrid` : array with parameters values
* `Start` : vector of areas from which agents randomly pick starting point
* `End` : vector of areas from which agents randomly pick ending point
* `density_factor` : road length reserved for one vehicle
* `K_Paths_Dict` : dictionary with multiple shortest paths (values) between vertices (keys)
* `mapdata` : OpenStreetMapX MapData object with road network data
"""
function run_parameter_analysis(GridElement::Int,
                                ParamGrid::Vector{Tuple{Int,Float64,Int,Int,Float64,Int,Float64}},
                                Start::Vector{Rect}, End::Vector{Rect}, density_factor::Float64,
                                K_Paths_Dict::Dict{Tuple{Int,Int},Array{Vector{Int}}},
                                mapdata::OpenStreetMapX.MapData)
  startTime = time()
  p = ParamGrid[GridElement]
  #Generating agents
  Agents = generate_agents(mapdata, p[3], Start, End, p[2], p[6], p[5], p[7], K_Paths_Dict)
  #Running base simulation - no V2I system
  BaseOutput = simulation_run(:base, mapdata, Agents, density_factor)
  #Running simulation with smart agents - V2I system enabled
  SmartOutput = simulation_run(:smart, mapdata, Agents, density_factor,
                              K_Paths_Dict, p[4], p[5], p[6], false)
  step_stat = gather_statistics(getfield.(Agents,:smart),
                                      BaseOutput.TravelTimes,
                                      SmartOutput.TravelTimes)
  runtime = time()-startTime
  print("$(GridElement),$(p[1]),$(p[2]),$(p[3]),$(p[4]),$(p[5]),$(p[6]),$(p[7]),")
  print("$(round(runtime,digits = 2)),$(step_stat[1]),$(step_stat[2]),$(step_stat[3]),")
  println("$(step_stat[4]),$(step_stat[5]),$(step_stat[6]),$(step_stat[7])")
end

"""
`run_parameter_analysis` run base and V2I simulation scenario to compare agents performance, returns comparison statistics

**Input parameters**
* `GridElement` : element from parameters grid to be used in simulation
* `ParamGrid` : array with parameters values
* `AgentsVec` : vector with agents used in simulation
* `density_factor` : road length reserved for one vehicle
* `K_Paths_Dict` : dictionary with multiple shortest paths (values) between vertices (keys)
* `mapdata` : OpenStreetMapX MapData object with road network data
"""
function run_parameter_analysis(GridElement::Int,
                                ParamGrid::Vector{Tuple{Int,Float64,Int,Int,Float64,Int,Float64}},
                                AgentsVec::Vector{Agent}, density_factor::Float64,
                                K_Paths_Dict::Dict{Tuple{Int,Int},Array{Vector{Int}}},
                                mapdata::OpenStreetMapX.MapData)
  startTime = time()
  p = ParamGrid[GridElement]
  #Sample N agents from supplied vector and adjust smart agents number
  Agents = deepcopy(sample(AgentsVec, p[3], replace = false))
  N_int= Int(ceil(p[3]*p[2]))
  smart_ind = [trues(N_int); falses(p[3]-N_int)]
  for i in 1:p[3] Agents[i].smart = smart_ind[i] end
  #Running base simulation - no V2I system
  BaseOutput = simulation_run(:base, mapdata, Agents, density_factor)
  #Running simulation with smart agents - V2I system enabled
  SmartOutput = simulation_run(:smart, mapdata, Agents, density_factor,
                              K_Paths_Dict, p[4], p[5], p[6], false)
  step_stat = gather_statistics(getfield.(Agents,:smart),
                                      BaseOutput.TravelTimes,
                                      SmartOutput.TravelTimes)
  runtime = time()-startTime
  print("$(GridElement),$(p[1]),$(p[2]),$(p[3]),$(p[4]),$(p[5]),$(p[6]),$(p[7]),")
  print("$(round(runtime,digits = 2)),$(step_stat[1]),$(step_stat[2]),$(step_stat[3]),")
  println("$(step_stat[4]),$(step_stat[5]),$(step_stat[6]),$(step_stat[7])")
end


# This file provides most of the implementation of multi-robot target *tracking*
# problems.
#
# The tracking problems involve estimation of target positions and optimization
# via maximizing mutual information.

using LinearAlgebra

export MultiRobotTargetTrackingConfigs, MultiRobotTargetTrackingProblem

# Automatically switch to sparse filters for large numbers of robots
const num_robots_sparse_filtering_threshold = 15
const sparsity_threshold=1e-3

# Stores configuration variables for multi-robot target tracking
struct MultiRobotTargetTrackingConfigs
  grid::Grid
  sensor::RangingSensor
  horizon::Int64

  # Configuration for the MCTS solver
  solver_iterations::Int64
  solver_information_samples::Int64

  # We will generally use a larger number of samples when querying actual
  # solution values
  objective_information_samples::Int64

  robot_target_range_limit::Float64

  function MultiRobotTargetTrackingConfigs(;
                                           grid,
                                           horizon=default_horizon,
                                           sensor=RangingSensor(),
                                           solver_iterations =
                                           default_num_iterations[horizon],
                                           solver_information_samples =
                                             default_solver_information_samples,
                                           objective_information_samples =
                                             default_num_information_samples,
                                           robot_target_range_limit = Inf
                                          )

    new(grid, sensor, horizon, solver_iterations, solver_information_samples,
        objective_information_samples, robot_target_range_limit)
  end
end

# Constructor that can infer parameters based on the number of robots
#
# Note that inferred parameters can still be overriden, but we probably don't
# want to do so
function MultiRobotTargetTrackingConfigs(num_robots;
                                         kwargs...)
  grid=Grid(num_robots=num_robots)
  MultiRobotTargetTrackingConfigs(;grid=grid, kwargs...)
end

# Copy constructor with overrides
function
  MultiRobotTargetTrackingConfigs(configs::MultiRobotTargetTrackingConfigs;
                                  kwargs...)

  fields = fieldnames(MultiRobotTargetTrackingConfigs)
  old_settings = Dict(field => getfield(configs, field) for field in fields)
  settings = merge(old_settings, kwargs)

  MultiRobotTargetTrackingConfigs(; settings...)
end

# Solution elements consist of the robot index and the associated trajectory
# We construct solutions as such because the output may not have the same
# ordering as the robots
#
# Warning: this "problem" object may become invalid if any of the underlying
# objects change and may copy some but not of the inputs.
struct MultiRobotTargetTrackingProblem{F<:AbstractFilter} <: AbstractMultiRobotProblem
  # Target tracking problems are defined by vectors of robot states
  partition_matroid::Vector{State}

  target_filters::Vector{F}
  filter_means::Vector{Tuple{Float64, Float64}}

  configs::MultiRobotTargetTrackingConfigs

  function MultiRobotTargetTrackingProblem(robot_states::Vector{State},
                 target_filters::Vector{F},
                 configs::MultiRobotTargetTrackingConfigs) where F <: AbstractFilter

    if length(robot_states) > num_robots_sparse_filtering_threshold
      sparse_filters = map(x->SparseFilter(x, threshold=sparsity_threshold),
                           target_filters)

      filter_means = compute_filter_means(sparse_filters)

      new{SparseFilter{Int64}}(robot_states, sparse_filters, filter_means,
                               configs)
    else
      filter_means = compute_filter_means(target_filters)

      new{F}(robot_states, target_filters, filter_means, configs)
    end
  end
end

# Construct a target tracking problem with configs. I may or may not need this
function MultiRobotTargetTrackingProblem(robot_states::Vector{State},
                                         target_filters::Vector{<:AbstractFilter};
                                         kwargs...)
  configs = MultiRobotTargetTrackingConfigs(;kwargs...)
  MultiRobotTargetTrackingProblem(robot_states, target_filters, configs)
end

function objective(p::MultiRobotTargetTrackingProblem, X)
  configs = p.configs

  trajectories = map(last, X)

  sum(p.target_filters) do filter
    finite_horizon_information(configs.grid, filter, configs.sensor,
                               trajectories;
                               num_samples=configs.objective_information_samples
                              ).reward
  end
end

# Returns distance from robot state to target mean
# (due to convexity and Jensen's inequality, this is a lower bound on the range)
function target_mean_distance(p::MultiRobotTargetTrackingProblem,
                     robot_state::State, target_index::Integer)
  norm(robot_state .- p.filter_means[target_index])
end
function target_mean_distance(p::MultiRobotTargetTrackingProblem,
                     robot_index::Integer, target_index::Integer)
  target_mean_distance(p, p.partition_matroid[robot_index], target_index)
end

function filter_targets_in_range(p::MultiRobotTargetTrackingProblem,
                                 robot_index::Integer;
                                 range_limit=p.configs.robot_target_range_limit
                                )

  # filter by range, producing tuples
  in_bounds = filter(1:length(p.filter_means)) do target_index
    target_mean_distance(p, robot_index, target_index) < range_limit
  end

  map(x->p.target_filters[x[1]], in_bounds)
end

function solve_block(p::MultiRobotTargetTrackingProblem, block::Integer,
                     selections::Vector)
  configs = p.configs

  trajectories = map(last, selections)

  # Select filters in range
  local targets_in_range
  if p.configs.robot_target_range_limit == Inf
    targets_in_range = p.target_filters
  else
    targets_in_range = filter_targets_in_range(p, block)
  end

  problem = SingleRobotTargetTrackingProblem(configs.grid, configs.sensor,
                                             configs.horizon, targets_in_range,
                                             prior_trajectories=trajectories,
                                             num_information_samples=
                                             configs.solver_information_samples
                                            )

  state = get_state(p, block)

  solution = solve_single_robot(problem, state,
                                n_iterations=configs.solver_iterations)

  (block, solution.trajectory)
end

function compute_filter_means(target_filters::Vector{<:AbstractFilter})
  map(target_filters) do x
    Tuple{Float64, Float64}(weighted_average(x))
  end
end


export
    LocalSimulation,
    simulate
"""
    LocalSimulation

Describe a Local-BPS sampling
"""
struct LocalSimulation
    fg::FactorGraph
    x0::Vector{AllowedVarType}  # Starting point
    v0::Vector{AllowedVarType}  # Starting velocity
    T::Real                     # Simulation time
    maxnevents::Int             # Maximum number of events
    lambdaref::Real             # Refreshment rate
    # -- implicit
    nvars::Int                  # Number of nodes
    nfactors::Int               # Number of factors
    function LocalSimulation(fg, x0, v0, T, maxnevents, lambdaref)
        # check none of the default named arguments went through
        cond =  !(  fg == :undefined ||
                    x0 == :undefined ||
                    v0 == :undefined ||
                    T  == :undefined )
        @assert cond "One or more essential argument undefined"
        cond = (length(x0) == length(v0) == fg.structure.nvars > 0)
        @assert cond "inconsistent dims"
        cond = (T > 0.0 && maxnevents > 0 && lambdaref > 0.0)
        @assert cond "inconsistent params"
        # creating the object
        new( fg, x0, v0, T,
             maxnevents,
             lambdaref,
             length(x0),
             length(fg.factors) )
    end
end
# Constructor with named arguments
function LocalSimulation(;
            factorgraph = :undefined,
            x0          = :undefined,
            v0          = :undefined,
            T           = :undefined,
            maxnevents  = 50000,
            lambdaref   = 0.5)::LocalSimulation
    # calling the unnamed constructor
    LocalSimulation( factorgraph,
                     x0, v0, T,
                     maxnevents,
                     lambdaref )
end

"""
    simulate(sim)

Run a local BPS simulation following the specifications of `sim`.
"""
function simulate(sim::LocalSimulation)::Tuple{AllEventList, Dict}

    (start, all_evlist, pq, tref) = ls_init(sim)

    # counters for the first branch (bounce) and the second branch (refresh)
    ev_firstbranch  = 0
    ev_secondbranch = 0
    nevents         = 0 # sum of above two
    # global clock
    globalclock     = 0.0

    # global iteration
    prog = Progress(sim.maxnevents, 1)

    for evnum in 1:sim.maxnevents
        # get bounce and dequeue
        (fidx, tbounce) = peek(pq)
        dequeue!(pq)
        t = tbounce
        globalclock += t
        # ------------
        # FIRST BRANCH (t<tref => )
        if t < tref
            ev_firstbranch += 1
            ls_firstbranch!(sim.fg, fidx, all_evlist, pq, t)
        # -------------
        # SECOND BRANCH (t>=tref => refreshment)
        else
            ev_secondbranch += 1
            pq    = ls_refreshment(sim.fg, tref, all_evlist)
            tref += randexp()/sim.lambdaref
        end
        # if use while loop
        # nevents += 1
        next!(prog)
    end

    details = Dict(
        "clocktime"       => time()-start,
        "nevents"         => nevents,
        "ev_firstbranch"  => ev_firstbranch,
        "ev_secondbranch" => ev_secondbranch,
        "globalclock"     => globalclock
    )

    (all_evlist, details)
end

# =============================================================================
# NOTE: The functions with a `ls_` prefix are helper functions exclusively
# designed for the local simulation. They aren't expected to be used elsewhere.

"""
    ls_init(sim)

Initialise a local simulation: store the starting clock time, create a list of
all eventlists and push the original event, create the priorityqueue and fill
it with the initial bounce times
"""
function ls_init(sim::LocalSimulation)
    # instantiate wall clock
    start = time()
    # initialisation of the eventlists for every node in the graph
    all_evlist = AllEventList([typeof(sim.x0[i]) for i in 1:sim.nvars])
    for i in 1:sim.nvars
        evi = Event(sim.x0[i], sim.v0[i], 0.0)
        pushevent!(all_evlist.evl[i], evi)
    end
    # initialisation of the priority queue and the refreshment time
    pq   = PriorityQueue{Int, Real}()
    tref = randexp()/sim.lambdaref
    # filling of the priority queue with initial position
    for fidx in 1:sim.fg.structure.nfactors
        (xf, vf, g) = ls_retrieve(sim.fg, fidx, all_evlist, 0.0)
        ls_updatepq!(pq, sim.fg, fidx, xf, vf, g, 0.0)
    end
    return (start, all_evlist, pq, tref)
end

"""
    ls_firstbranch(fg, fidx, all_evlist, pq, t)

Operation corresponding to the first branch of events in simulate.
"""
function ls_firstbranch!(fg::FactorGraph, fidx::Int, all_evlist::AllEventList,
                         pq::PriorityQueue, t::Real)
    # retrieve xf, vf corresponding to factor
    (xf, vf, g, vars) = ls_retrieve(fg, fidx, all_evlist, t, true)
    ls_saveupdate!(all_evlist, vars, xf, vf, t)
    ls_updatepq!(pq, fg, fidx, xf, vf, g, t)
    # same story for linked factors (fp)
    for fpidx in linkedfactors(fg, fidx)
        # we don't need to retrieve `vars` here
        (xfp, vfp, gp) = ls_retrieve(fg, fpidx, all_evlist, t)
        ls_updatepq!(pq, fg, fpidx, xfp, vfp, gp, t)
    end
    return (all_evlist, pq)
end

"""
    ls_reshape(u, v)

(Local Simulation, helper function) Reshape a single vector `u` to a vector of
AllowedVarType following the structure of `v`. For example, let
`v = [0., [0., 0.]]` and `u=[1., 2., 3.]` then `ls_reshape(u,v)` will return
`[1., [2., 3.]]`.
"""
function ls_reshape(u::Vector{<:Real}, v::V) where V <: Vector{AllowedVarType}
    w = similar(v)
    # offsets to know where to look for next block of information
    offset = 0
    for i ∈ 1:length(v)
        tmpi = length(v[i])
        # if length is 1, return only single point, otherwise vector
        @inbounds w[i] = (tmpi==1) ? u[offset+1] : u[offset .+ (1:tmpi)]
        offset += tmpi
    end
    return w
end

"""
    ls_retrieve(fg, fidx, all_evlist, t, doreflect)

(Local simulation, helper function) Retrieve the structure corresponding to all
positions of the nodes attached to a factor (`xf`) and the velocities (`vf`) as
well as the gradient at `xf` and the list of indices of the attached variables.
"""
function ls_retrieve(fg::FactorGraph, fidx::Int,
                    all_evlist::AllEventList,
                    t::Real, doreflect::Bool=false)
    # indices of the variable associated with factor fidx
    vars = assocvariables(fg, fidx)
    # allocate xf, vf, note the different variables
    # don't necessarily have the same types
    vf = Vector{AllowedVarType}(undef, length(vars))
    xf = similar(vf)
    # shortcut to initial event
    if t == 0.0
        for (i, k) ∈ enumerate(vars)
            ev = getevent(all_evlist.evl[k], 1)
            xf[i], vf[i] = ev.x, ev.v
        end
    # retrieving later than initial event (standard case)
    else
        # retrieve vf, xf(t)
        for (i, k) ∈ enumerate(vars)
            # get the eventlist corresponding to variable k
            ev = getlastevent(all_evlist.evl[k])
            # store the components and extrapolate from ev.x as needed
            vf[i] = ev.v
            xf[i] = ev.x + (t-ev.t)*ev.v
        end
    end
    # compute the gradient of the factor's log likelihood at xf
    g  = fg.factors[fidx].gll(vcat(xf...))
    # compute the reflection (if required)
    vf = doreflect ? ls_reshape(reflect_bps!(g, vcat(vf...)), vf) : vf
    (xf, vf, g, vars)
end

"""
    ls_saveupdate!(all_evlist, vars, xf, vf, t)

(Local simulation, helper function) Append the information contained in `xf`,
`vf` in the appropriate EventLists corresponding to the relevant variables
indexed in `vars` (variables associated with the factor `f`, `xf` and `vf` are
the recovered positions and velocities obtained when considering factor `f`).
"""
function ls_saveupdate!(all_evlist::AllEventList, vars::Vector{Int},
                        xf::Vector{AllowedVarType}, vf::Vector{AllowedVarType},
                        t::Real)
    # Add the new sample to the list
    for (i, k) ∈ enumerate(vars)
        pushevent!(all_evlist.evl[k], Event(xf[i], vf[i], t))
    end
    return all_evlist
end

"""
    ls_updatepq!(pq, fg, fidx, xf, vf, g, t)

(Local simulation, helper function) Compute the bounce time (possibly via
thinning depending on the factor's implementation) associated with a factor and
push it to the PriorityQueue `pq`.
"""
function ls_updatepq!(pq::PriorityQueue, fg::FactorGraph, fidx::Int,
                      xf::Vector{AllowedVarType}, vf::Vector{AllowedVarType},
                      g::Vector{<:Real}, t::Real)
    # useful temporary variables
    vcxf, vcvf = vcat(xf...), vcat(vf...)
    # Update time in Priority Queue for the current factor
    acc  = false
    tauf = 0.0
    while !acc
        bounce = fg.factors[fidx].nextevent(vcxf, vcvf)
        acc, tauf = bounce.dobounce(g, vcvf), bounce.tau
    end
    # add the time to current priorityQueue
    pq[fidx] = t+tauf
    return pq
end

"""
    ls_random(evl)

(Local simulation, helper function) Generate N(0,1) random numbers of dimension
corresponding to the passed `EventList`.
"""
function ls_random(evl::EventList)
    l = length(evl.xs[end])
    return l > 1 ? randn(l) : randn()
end

"""
    ls_refreshment(fg, t, all_evlist)

(Local simulation, helper function) Generate a new PriorityQueue by drawing
new velocities for each node and storing it (cf. NewQueue algorithm in the BPS
paper).
"""
function ls_refreshment(fg::FactorGraph, t::Real, all_evlist::AllEventList)
    # draw a new bunch of velocities for each node
    v = Vector{AllowedVarType}(undef, fg.structure.nvars)
    for i ∈ 1:fg.structure.nvars
        @inbounds v[i] = ls_random(all_evlist.evl[i])
    end
    # Instantiate a new priority queue
    pq = PriorityQueue{Int, Real}()
    for fidx ∈ 1:fg.structure.nfactors
        # retrieve xf, vf corresponding to factor
        (xf, vf, g, vars) = ls_retrieve(fg, fidx, all_evlist, t)
        # update the priority queue, note v[vars] and not vf
        # since we are using the refreshed velocity
        ls_updatepq!(pq, fg, fidx, xf, v[vars], g, t)
    end
    return pq
end


"""
    Channel

Struct containing asymptotic kernels for a channel.
* `q1   :: Matrix{Float64}`   : kernel with both fermionic frequencies -> Inf
* `q2_1 :: Array{Float64, 3}` : kernel with second fermionic frequency -> Inf
* `q2_2 :: Array{Float64, 3}` : kernel with first fermionic frequency -> Inf
* `q3   :: Array{Float64, 4}` : full channel
"""
struct Channel
    q1   :: Matrix{Float64}
    q2_1 :: Array{Float64, 3}
    q2_2 :: Array{Float64, 3}
    q3   :: Array{Float64, 4}
end

# generate channel dummy
function get_channel_empty(
    r :: Reduced_lattice,
    m :: Mesh
    ) :: Channel

    num_sites = length(r.sites)

    # init kernels
    q1   = zeros(Float64, num_sites, m.num_Ω)
    q2_1 = zeros(Float64, num_sites, m.num_Ω, m.num_ν)
    q2_2 = zeros(Float64, num_sites, m.num_Ω, m.num_ν)
    q3   = zeros(Float64, num_sites, m.num_Ω, m.num_ν, m.num_ν)

    # build channel
    ch = Channel(q1, q2_1, q2_2, q3)

    return ch
end

# get interpolated value of q1
function get_q1(
    site :: Int64,
    p    :: Param,
    ch   :: Channel
    )    :: Float64

    val  = p.lower_weight * ch.q1[site, p.lower_index]
    val += p.upper_weight * ch.q1[site, p.upper_index]

    return val
end

# get interpolated value of q1 on all lattice sites
function get_q1_avx!(
    r        :: Reduced_lattice,
    p        :: Param,
    ch       :: Channel,
    temp     :: SubArray{Float64, 1, Array{Float64, 3}},
    exchange :: Bool,
    sgn      :: Float64
    )        :: Nothing

    # deref channel
    q1 = ch.q1

    # deref param
    lower_index  = p.lower_index
    upper_index  = p.upper_index
    lower_weight = p.lower_weight
    upper_weight = p.upper_weight

    # check for site exchange
    if exchange
        indices = r.exchange

        @turbo unroll = 1 for i in eachindex(temp)
            val      = lower_weight * q1[indices[i], lower_index]
            val     += upper_weight * q1[indices[i], upper_index]
            temp[i] += sgn * val
        end
    else
        @turbo unroll = 1 for i in eachindex(temp)
            val      = lower_weight * q1[i, lower_index]
            val     += upper_weight * q1[i, upper_index]
            temp[i] += sgn * val
        end
    end

    return nothing
end

# get interpolated value of q2_1
function get_q2_1(
    site :: Int64,
    p1   :: Param,
    p2   :: Param,
    ch   :: Channel
    )    :: Float64

    val  = p1.lower_weight * p2.lower_weight * ch.q2_1[site, p1.lower_index, p2.lower_index]
    val += p1.upper_weight * p2.lower_weight * ch.q2_1[site, p1.upper_index, p2.lower_index]
    val += p1.lower_weight * p2.upper_weight * ch.q2_1[site, p1.lower_index, p2.upper_index]
    val += p1.upper_weight * p2.upper_weight * ch.q2_1[site, p1.upper_index, p2.upper_index]

    return val
end

# get interpolated value of q2_1 on all lattice sites
function get_q2_1_avx!(
    r        :: Reduced_lattice,
    p1       :: Param,
    p2       :: Param,
    ch       :: Channel,
    temp     :: SubArray{Float64, 1, Array{Float64, 3}},
    exchange :: Bool,
    sgn      :: Float64
    )        :: Nothing

    # deref channel
    q2_1 = ch.q2_1

    # deref param
    lower_index1  = p1.lower_index
    upper_index1  = p1.upper_index
    lower_weight1 = p1.lower_weight
    upper_weight1 = p1.upper_weight

    lower_index2  = p2.lower_index
    upper_index2  = p2.upper_index
    lower_weight2 = p2.lower_weight
    upper_weight2 = p2.upper_weight

    # check for site exchange
    if exchange
        indices = r.exchange

        @turbo unroll = 1 for i in eachindex(temp)
            val      = lower_weight1 * lower_weight2 * q2_1[indices[i], lower_index1, lower_index2]
            val     += upper_weight1 * lower_weight2 * q2_1[indices[i], upper_index1, lower_index2]
            val     += lower_weight1 * upper_weight2 * q2_1[indices[i], lower_index1, upper_index2]
            val     += upper_weight1 * upper_weight2 * q2_1[indices[i], upper_index1, upper_index2]
            temp[i] += sgn * val
        end
    else
        @turbo unroll = 1 for i in eachindex(temp)
            val      = lower_weight1 * lower_weight2 * q2_1[i, lower_index1, lower_index2]
            val     += upper_weight1 * lower_weight2 * q2_1[i, upper_index1, lower_index2]
            val     += lower_weight1 * upper_weight2 * q2_1[i, lower_index1, upper_index2]
            val     += upper_weight1 * upper_weight2 * q2_1[i, upper_index1, upper_index2]
            temp[i] += sgn * val
        end
    end

    return nothing
end

# get interpolated value of q2_2
function get_q2_2(
    site :: Int64,
    p1   :: Param,
    p2   :: Param,
    ch   :: Channel
    )    :: Float64

    val  = p1.lower_weight * p2.lower_weight * ch.q2_2[site, p1.lower_index, p2.lower_index]
    val += p1.upper_weight * p2.lower_weight * ch.q2_2[site, p1.upper_index, p2.lower_index]
    val += p1.lower_weight * p2.upper_weight * ch.q2_2[site, p1.lower_index, p2.upper_index]
    val += p1.upper_weight * p2.upper_weight * ch.q2_2[site, p1.upper_index, p2.upper_index]

    return val
end

# get interpolated value of q2_2 on all lattice sites
function get_q2_2_avx!(
    r        :: Reduced_lattice,
    p1       :: Param,
    p2       :: Param,
    ch       :: Channel,
    temp     :: SubArray{Float64, 1, Array{Float64, 3}},
    exchange :: Bool,
    sgn      :: Float64
    )        :: Nothing

    # deref channel
    q2_2 = ch.q2_2

    # deref param
    lower_index1  = p1.lower_index
    upper_index1  = p1.upper_index
    lower_weight1 = p1.lower_weight
    upper_weight1 = p1.upper_weight

    lower_index2  = p2.lower_index
    upper_index2  = p2.upper_index
    lower_weight2 = p2.lower_weight
    upper_weight2 = p2.upper_weight

    # check for site exchange
    if exchange
        indices = r.exchange

        @turbo unroll = 1 for i in eachindex(temp)
            val      = lower_weight1 * lower_weight2 * q2_2[indices[i], lower_index1, lower_index2]
            val     += upper_weight1 * lower_weight2 * q2_2[indices[i], upper_index1, lower_index2]
            val     += lower_weight1 * upper_weight2 * q2_2[indices[i], lower_index1, upper_index2]
            val     += upper_weight1 * upper_weight2 * q2_2[indices[i], upper_index1, upper_index2]
            temp[i] += sgn * val
        end
    else
        @turbo unroll = 1 for i in eachindex(temp)
            val      = lower_weight1 * lower_weight2 * q2_2[i, lower_index1, lower_index2]
            val     += upper_weight1 * lower_weight2 * q2_2[i, upper_index1, lower_index2]
            val     += lower_weight1 * upper_weight2 * q2_2[i, lower_index1, upper_index2]
            val     += upper_weight1 * upper_weight2 * q2_2[i, upper_index1, upper_index2]
            temp[i] += sgn * val
        end
    end

    return nothing
end

# get interpolated value of q3
function get_q3(
    site :: Int64,
    p1   :: Param,
    p2   :: Param,
    p3   :: Param,
    ch   :: Channel
    )    :: Float64

    val  = p1.lower_weight * p2.lower_weight * p3.lower_weight * ch.q3[site, p1.lower_index, p2.lower_index, p3.lower_index]
    val += p1.upper_weight * p2.lower_weight * p3.lower_weight * ch.q3[site, p1.upper_index, p2.lower_index, p3.lower_index]
    val += p1.lower_weight * p2.upper_weight * p3.lower_weight * ch.q3[site, p1.lower_index, p2.upper_index, p3.lower_index]
    val += p1.upper_weight * p2.upper_weight * p3.lower_weight * ch.q3[site, p1.upper_index, p2.upper_index, p3.lower_index]
    val += p1.lower_weight * p2.lower_weight * p3.upper_weight * ch.q3[site, p1.lower_index, p2.lower_index, p3.upper_index]
    val += p1.upper_weight * p2.lower_weight * p3.upper_weight * ch.q3[site, p1.upper_index, p2.lower_index, p3.upper_index]
    val += p1.lower_weight * p2.upper_weight * p3.upper_weight * ch.q3[site, p1.lower_index, p2.upper_index, p3.upper_index]
    val += p1.upper_weight * p2.upper_weight * p3.upper_weight * ch.q3[site, p1.upper_index, p2.upper_index, p3.upper_index]

    return val
end

# get interpolated value of q3 on all lattice sites
function get_q3_avx!(
    r        :: Reduced_lattice,
    p1       :: Param,
    p2       :: Param,
    p3       :: Param,
    ch       :: Channel,
    temp     :: SubArray{Float64, 1, Array{Float64, 3}},
    exchange :: Bool,
    sgn      :: Float64
    )        :: Nothing

    # deref channel
    q3 = ch.q3

    # deref param
    lower_index1  = p1.lower_index
    upper_index1  = p1.upper_index
    lower_weight1 = p1.lower_weight
    upper_weight1 = p1.upper_weight

    lower_index2  = p2.lower_index
    upper_index2  = p2.upper_index
    lower_weight2 = p2.lower_weight
    upper_weight2 = p2.upper_weight

    lower_index3  = p3.lower_index
    upper_index3  = p3.upper_index
    lower_weight3 = p3.lower_weight
    upper_weight3 = p3.upper_weight

    # check for site exchange
    if exchange
        indices = r.exchange

        @turbo unroll = 1 for i in eachindex(temp)
            val      = lower_weight1 * lower_weight2 * lower_weight3 * q3[indices[i], lower_index1, lower_index2, lower_index3]
            val     += upper_weight1 * lower_weight2 * lower_weight3 * q3[indices[i], upper_index1, lower_index2, lower_index3]
            val     += lower_weight1 * upper_weight2 * lower_weight3 * q3[indices[i], lower_index1, upper_index2, lower_index3]
            val     += upper_weight1 * upper_weight2 * lower_weight3 * q3[indices[i], upper_index1, upper_index2, lower_index3]
            val     += lower_weight1 * lower_weight2 * upper_weight3 * q3[indices[i], lower_index1, lower_index2, upper_index3]
            val     += upper_weight1 * lower_weight2 * upper_weight3 * q3[indices[i], upper_index1, lower_index2, upper_index3]
            val     += lower_weight1 * upper_weight2 * upper_weight3 * q3[indices[i], lower_index1, upper_index2, upper_index3]
            val     += upper_weight1 * upper_weight2 * upper_weight3 * q3[indices[i], upper_index1, upper_index2, upper_index3]
            temp[i] += sgn * val
        end
    else
        @turbo unroll = 1 for i in eachindex(temp)
            val      = lower_weight1 * lower_weight2 * lower_weight3 * q3[i, lower_index1, lower_index2, lower_index3]
            val     += upper_weight1 * lower_weight2 * lower_weight3 * q3[i, upper_index1, lower_index2, lower_index3]
            val     += lower_weight1 * upper_weight2 * lower_weight3 * q3[i, lower_index1, upper_index2, lower_index3]
            val     += upper_weight1 * upper_weight2 * lower_weight3 * q3[i, upper_index1, upper_index2, lower_index3]
            val     += lower_weight1 * lower_weight2 * upper_weight3 * q3[i, lower_index1, lower_index2, upper_index3]
            val     += upper_weight1 * lower_weight2 * upper_weight3 * q3[i, upper_index1, lower_index2, upper_index3]
            val     += lower_weight1 * upper_weight2 * upper_weight3 * q3[i, lower_index1, upper_index2, upper_index3]
            val     += upper_weight1 * upper_weight2 * upper_weight3 * q3[i, upper_index1, upper_index2, upper_index3]
            temp[i] += sgn * val
        end
    end

    return nothing
end

# get interpolated value of channel for a given frequency buffer
function get_channel(
    site :: Int64,
    b    :: Buffer,
    ch   :: Channel
    )    :: Float64

    val = 0.0

    if b.kernel == 1
        val = get_q1(site, b.p1, ch)
    elseif b.kernel == 2
        val = get_q2_1(site, b.p1, b.p2, ch)
    elseif b.kernel == 3
        val = get_q2_2(site, b.p1, b.p3, ch)
    elseif b.kernel == 4
        val = get_q3(site, b.p1, b.p2, b.p3, ch)
    end

    return val
end

# get interpolated value of channel for a given frequency buffer on all lattice sites
function get_channel_avx!(
    r        :: Reduced_lattice,
    b        :: Buffer,
    ch       :: Channel,
    temp     :: SubArray{Float64, 1, Array{Float64, 3}},
    exchange :: Bool,
    sgn      :: Float64
    )        :: Nothing

    if b.kernel == 1
        get_q1_avx!(r, b.p1, ch, temp, exchange, sgn)
    elseif b.kernel == 2
        get_q2_1_avx!(r, b.p1, b.p2, ch, temp, exchange, sgn)
    elseif b.kernel == 3
        get_q2_2_avx!(r, b.p1, b.p3, ch, temp, exchange, sgn)
    elseif b.kernel == 4
        get_q3_avx!(r, b.p1, b.p2, b.p3, ch, temp, exchange, sgn)
    end

    return nothing
end

# replace channel with another channel
function replace_with!(
    ch1 :: Channel,
    ch2 :: Channel
    )   :: Nothing

    @turbo ch1.q1   .= ch2.q1
    @turbo ch1.q2_1 .= ch2.q2_1
    @turbo ch1.q2_2 .= ch2.q2_2
    @turbo ch1.q3   .= ch2.q3

    return nothing
end

# multiply channel with factor
function mult_with!(
    ch  :: Channel,
    fac :: Float64
    )   :: Nothing

    @turbo ch.q1   .*= fac
    @turbo ch.q2_1 .*= fac
    @turbo ch.q2_2 .*= fac
    @turbo ch.q3   .*= fac

    return nothing
end

# multiply channel with factor and add to other channel
function mult_with_add_to!(
    ch2 :: Channel,
    fac :: Float64,
    ch1 :: Channel
    )   :: Nothing

    @turbo ch1.q1   .+= fac .* ch2.q1
    @turbo ch1.q2_1 .+= fac .* ch2.q2_1
    @turbo ch1.q2_2 .+= fac .* ch2.q2_2
    @turbo ch1.q3   .+= fac .* ch2.q3

    return nothing
end

"""
    get_abs_max(
        ch :: Channel
        )  :: Float64

Returns maximum absolute value of a channel.
"""
function get_abs_max(
    ch :: Channel
    )  :: Float64

    max_ch = maximum(abs.(ch.q3))

    return max_ch
end

# resample a channel to new meshes via trilinear interpolation
function resample_from_to!(
    Ω_old  :: Vector{Float64},
    ν_old  :: Vector{Float64},
    ch_old :: Channel,
    Ω_new  :: Vector{Float64},
    ν_new  :: Vector{Float64},
    ch_new :: Channel
    )      :: Nothing

    # get dimensions
    num_sites = size(ch_old.q2_1, 1)
    num_Ω     = size(ch_old.q2_1, 2)
    num_ν     = size(ch_old.q2_1, 3)

    # resample q1
    for w in 1 : num_Ω
        # get interpolation parameters
        p1 = get_param(Ω_new[w], Ω_old)

        for site in 1 : num_sites
            ch_new.q1[site, w] = get_q1(site, p1, ch_old)
        end
    end

    # resample q2_1 and q2_2
    for v in 1 : num_ν
        for w in 1 : num_Ω
            # get interpolation parameters
            p1 = get_param(Ω_new[w], Ω_old)
            p2 = get_param(ν_new[v], ν_old)

            for site in 1 : num_sites
                ch_new.q2_1[site, w, v] = get_q2_1(site, p1, p2, ch_old)
                ch_new.q2_2[site, w, v] = get_q2_2(site, p1, p2, ch_old)
            end
        end
    end

    # resample q3
    for vp in 1 : num_ν
        for v in 1 : num_ν
            for w in 1 : num_Ω
                # get interpolation parameters
                p1 = get_param(Ω_new[w],  Ω_old)
                p2 = get_param(ν_new[v],  ν_old)
                p3 = get_param(ν_new[vp], ν_old)

                for site in 1 : num_sites
                    ch_new.q3[site, w, v, vp] = get_q3(site, p1, p2, p3, ch_old)
                end
            end
        end
    end

    return nothing
end


"""
    LensAssembly{T<:Real}

Structure which contains the elements of the optical system, these can be [`CSGTree`](@ref) or [`Surface`](@ref) objects.

In order to prevent type ambiguities bespoke structs are created for each possible number of elements e.g. `LensAssembly3`.
These are parameterized by the types of the elements to prevent ambiguities.
Basic surface types such as [`Rectangle`](@ref) (which can occur in large numbers) are stored independently in `Vector`s, so type paramters are only needed for CSG objects.

Each struct looks like this:
```julia
struct LensAssemblyN{T,T1,T2,...,TN} <: LensAssembly{T}
    axis::SVector{3,T}
    rectangles::Vector{Rectangle{T}}
    ellipses::Vector{Ellipse{T}}
    hexagons::Vector{Hexagon{T}}
    paraxials::Vector{ParaxialLens{T}}
    E1::T1
    E2::T2
    ...
    EN::TN
end
```

Where `Ti <: Union{Surface{T},CSGTree{T}}`.

To create a LensAssembly object the following functions can be used:
```julia
LensAssembly(elements::Vararg{Union{Surface{T},CSGTree{T},LensAssembly{T}}}; axis = SVector(0.0, 0.0, 1.0)) where {T<:Real}
```
"""
abstract type LensAssembly{T<:Real} end
export LensAssembly

Base.show(io::IO, ::Type{<:LensAssembly}) = print(io, "LensAssembly")
Base.show(io::IO, a::LensAssembly{T}) where {T<:Real} = print(io, "LensAssembly{$T}($(length(a.rectangles) + length(a.ellipses) + length(a.hexagons) + length(typed_elements(a))) elements)")

# in order to prevent type ambiguities lens assembly must be paramaterised by its elements, but we want lens assembly to have an arbitrary number of elements
# as such we have to use macros to create lens assembly structs of a given size, and methods to handle them in a type unambiguous way
# we may often end up with lots of simple surfaces (Rectangle, Ellipse) in the assembly, and we can store these more efficiently in Vectors
macro lensassembly_constructor(N)
    # generate the type definition for a lens assembly of size N
    typesig = [:($(Symbol("T$(i)")) <: Union{Surface{T},CSGTree{T}}) for i in 1:N]
    fields = [:($(Symbol("E$(i)"))::$(Symbol("T$(i)"))) for i in 1:N]
    name = Symbol("LensAssembly$N")
    return esc(quote
        struct $(name){T,$(typesig...)} <: LensAssembly{T}
            axis::SVector{3,T}
            rectangles::Vector{Rectangle{T}}
            ellipses::Vector{Ellipse{T}}
            hexagons::Vector{Hexagon{T}}
            paraxials::Vector{ParaxialLens{T}}
            $(fields...)
        end
    end)
end

# This is the basic idea of what each closestintersection implementation does:
# function closestintersection(ass::LensAssembly{T}, r::AbstractRay{T,N})::Union{Nothing,Intersection{T,N}} where {T<:Real,N}
#     αmin = typemax(T)
#     closest = nothing
#     for element in elements(ass)
#         allintscts = surfaceintersection(element, r)
#         if !(allintscts isa EmptyInterval)
#             intsct = closestintersection(allintscts)
#             if intsct !== nothing
#                 αcurr = α(intsct)
#                 if αcurr < αmin
#                     αmin = αcurr
#                     closest = intsct
#                 end
#             end
#         end
#     end
#     return closest
# end

macro lensassembly_intersection(N)
    # generates the _closestintersection() function for a lens assembly of size N, this should never be called directly, instead call closestintersection() in all cases
    type = :($(Symbol("LensAssembly$(N)")))
    fieldnames = [Symbol("E$(i)") for i in 1:N]
    checks = [quote
        intvl = surfaceintersection(obj.$(fieldnames[i]), r)
        if !(intvl isa EmptyInterval)
            intsct = closestintersection(intvl::Union{Interval{T},DisjointUnion{T}})
            if intsct !== nothing
                αcurr = α(intsct)
                if αcurr < αmin
                    αmin = αcurr
                    closest = intsct
                end
            end
        end
    end for i in 1:N]
    return esc(quote
        function _closestintersection(obj::$(type){T}, r::AbstractRay{T,N})::Union{Nothing,Intersection{T,N}} where {T<:Real,N}
            αmin = typemax(T)
            closest = nothing
            for rect in obj.rectangles
                intvl = surfaceintersection(rect, r)
                if !(intvl isa EmptyInterval)
                    intsct = closestintersection(intvl::Union{Interval{T},DisjointUnion{T}})
                    if intsct !== nothing
                        αcurr = α(intsct)
                        if αcurr < αmin
                            αmin = αcurr
                            closest = intsct
                        end
                    end
                end
            end
            for ell in obj.ellipses
                intvl = surfaceintersection(ell, r)
                if !(intvl isa EmptyInterval)
                    intsct = closestintersection(intvl::Union{Interval{T},DisjointUnion{T}})
                    if intsct !== nothing
                        αcurr = α(intsct)
                        if αcurr < αmin
                            αmin = αcurr
                            closest = intsct
                        end
                    end
                end
            end
            for hex in obj.hexagons
                intvl = surfaceintersection(hex, r)
                if !(intvl isa EmptyInterval)
                    intsct = closestintersection(intvl::Union{Interval{T},DisjointUnion{T}})
                    if intsct !== nothing
                        αcurr = α(intsct)
                        if αcurr < αmin
                            αmin = αcurr
                            closest = intsct
                        end
                    end
                end
            end
            for par in obj.paraxials
                intvl = surfaceintersection(par, r)
                if !(intvl isa EmptyInterval)
                    intsct = closestintersection(intvl::Union{Interval{T},DisjointUnion{T}})
                    if intsct !== nothing
                        αcurr = α(intsct)
                        if αcurr < αmin
                            αmin = αcurr
                            closest = intsct
                        end
                    end
                end
            end
            $(checks...)
            return closest
        end
    end)
end

macro lensassembly_elements(N)
    # generates the _elements() function for a lens assembly of size N, this should never be called directly, instead call elements() in all cases
    type = :($(Symbol("LensAssembly$(N)")))
    fieldnames = [Symbol("E$(i)") for i in 1:N]
    tuple = [:(obj.$(fieldnames[i])) for i in 1:N]
    return esc(quote
        function _typed_elements(obj::$(type){T}) where {T<:Real}
            return tuple($(tuple...))
        end
    end)
end

# pregenerate a small number at compile time
for N in 1:PREGENERATED_LENS_ASSEMBLY_SIZE
    type = :($(Symbol("LensAssembly$(N)")))
    if !isdefined(OpticSim, type)
        @eval @lensassembly_constructor($N)
        @eval @lensassembly_intersection($N)
        @eval @lensassembly_elements($N)
    end
end

function LensAssembly(elements::Vararg{Union{Surface{T},CSGTree{T},LensAssembly{T}}}; axis::SVector{3,T} = SVector{3,T}(0.0, 0.0, 1.0)) where {T<:Real}
    # make the actual object
    actual_elements = []
    rectangles = Vector{Rectangle{T}}(undef, 0)
    ellipses = Vector{Ellipse{T}}(undef, 0)
    hexagons = Vector{Hexagon{T}}(undef, 0)
    paraxials = Vector{ParaxialLens{T}}(undef, 0)
    for e in elements
        if e isa Rectangle{T}
            push!(rectangles, e)
        elseif e isa Ellipse{T}
            push!(ellipses, e)
        elseif e isa Hexagon{T}
            push!(hexagons, e)
        elseif e isa ParaxialLens{T}
            push!(paraxials, e)
        elseif e isa Surface{T} || e isa CSGTree{T}
            push!(actual_elements, e)
        else
            # unpack the lens assembly
            rectangles = vcat(rectangles, e.rectangles)
            ellipses = vcat(ellipses, e.ellipses)
            hexagons = vcat(hexagons, e.hexagons)
            for lae in OpticSim.typed_elements(e)
                push!(actual_elements, lae)
            end
        end
    end
    N = length(actual_elements)
    # make the methods for this N
    type = :($(Symbol("LensAssembly$(N)")))
    if !isdefined(OpticSim, type)
        @eval @lensassembly_constructor($N)
        @eval @lensassembly_intersection($N)
        @eval @lensassembly_elements($N)
    end
    eltypes = typeof.(actual_elements)
    naxis = normalize(axis)
    return eval(:($(type){$T,$(eltypes...)}($naxis, $rectangles, $ellipses, $hexagons, $paraxials, $(actual_elements...))))
end

function typed_elements(ass::LensAssembly{T}) where {T<:Real}
    # under certain circumstances we run into the world age problem (see https://discourse.julialang.org/t/how-to-bypass-the-world-age-problem/7012)
    # to avoid this we need to use invokelatest(), but this is type ambiguous and ruins performance, and in most cases isn't necessary
    # so we can use try/catch to only call it when we need (and warn appropriately), retaining performance in almost all cases
    # NEVER call _typed_elements directly, always use this method
    try
        _typed_elements(ass)
    catch MethodError
        @warn "First time call to LensAssembly of previously unused size, performance will be worse than normal." maxlog = 1
        Base.invokelatest(_typed_elements, ass)
    end
end

elements(ass::LensAssembly{T}) where {T<:Real} = (ass.rectangles..., ass.ellipses..., ass.hexagons..., ass.paraxials..., typed_elements(ass)...)
export elements

function closestintersection(ass::LensAssembly{T}, r::AbstractRay{T,N})::Union{Nothing,Intersection{T,N}} where {T<:Real,N}
    # under certain circumstances we run into the world age problem (see https://discourse.julialang.org/t/how-to-bypass-the-world-age-problem/7012)
    # to avoid this we need to use invokelatest(), but this is type ambiguous and ruins performance, and in most cases isn't necessary
    # so we can use try/catch to only call it when we need (and warn appropriately), retaining performance in almost all cases
    # NEVER call _closestintersection directly, always use this method
    try
        _closestintersection(ass, r)
    catch MethodError
        @warn "First time call to LensAssembly of previously unused size, performance will be worse than normal." maxlog = 1
        Base.invokelatest(_closestintersection, ass, r)
    end
end

function BoundingBox(la::LensAssembly{T}) where {T<:Real}
    es = elements(la)
    bbox = BoundingBox(es[1])
    for b in es[2:end]
        if b isa ParametricSurface{T} || b isa CSGTree{T}
            bbox = union(bbox, b)
        end
    end
    return bbox
end

axis(a::LensAssembly{T}) where {T<:Real} = a.axis
export axis

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

"""
    LensTrace{T<:Real,N}

Contains an intersection point and the ray segment leading to it from within an optical trace.
The ray carries the path length, power, wavelength, number of intersections and source number, all of which are accessible directly on this class too.

Has the following accessor methods:
```julia
ray(a::LensTrace{T,N}) -> OpticalRay{T,N}
intersection(a::LensTrace{T,N}) -> Intersection{T,N}
power(a::LensTrace{T,N}) -> T
wavelength(a::LensTrace{T,N}) -> T
pathlength(a::LensTrace{T,N}) -> T
point(a::LensTrace{T,N}) -> SVector{N,T}
uv(a::LensTrace{T,N}) -> SVector{2,T}
sourcenum(a::LensTrace{T,N}) -> Int
nhits(a::LensTrace{T,N}) -> Int
```
"""
struct LensTrace{T<:Real,N}
    ray::OpticalRay{T,N}
    intersection::Intersection{T,N}
end
export LensTrace

ray(a::LensTrace{T,N}) where {T<:Real,N} = a.ray
intersection(a::LensTrace{T,N}) where {T<:Real,N} = a.intersection
power(a::LensTrace{T,N}) where {T<:Real,N} = power(ray(a))
wavelength(a::LensTrace{T,N}) where {T<:Real,N} = wavelength(ray(a))
pathlength(a::LensTrace{T,N}) where {T<:Real,N} = pathlength(ray(a))
point(a::LensTrace{T,N}) where {T<:Real,N} = point(intersection(a))
uv(a::LensTrace{T,N}) where {T<:Real,N} = uv(intersection(a))
sourcenum(a::LensTrace{T,N}) where {T<:Real,N} = sourcenum(ray(a))
nhits(a::LensTrace{T,N}) where {T<:Real,N} = nhits(ray(a))
export intersection

function Base.print(io::IO, a::LensTrace{T,N}) where {T,N}
    println(io, "Ray\n$(ray(a))")
    println(io, "Lensintersection\n$(intersection(a))")
end

struct NoPower end
const nopower = NoPower()

# isnopower(::NoPower) = true
# isnopower(::Any) = false

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

"""
    trace(assembly::LensAssembly{T}, r::OpticalRay{T}, temperature::T = 20.0, pressure::T = 1.0; trackrays = nothing, test = false)

Returns the ray as it exits the assembly in the form of a [`LensTrace`](@ref) object if it hits any element in the assembly, otherwise `nothing`.
Recursive rays are offset by a small amount (`RAY_OFFSET`) to prevent it from immediately reintersecting the same lens element.

`trackrays` can be passed an empty vector to accumulate the `LensTrace` objects at each intersection of `ray` with a surface in the assembly.
"""
function trace(assembly::LensAssembly{T}, r::OpticalRay{T,N}, temperature::T = T(OpticSim.GlassCat.TEMP_REF), pressure::T = T(OpticSim.GlassCat.PRESSURE_REF); trackrays::Union{Nothing,Vector{LensTrace{T,N}}} = nothing, test::Bool = false, recursion::Int = 0)::Union{Nothing,NoPower,LensTrace{T,N}} where {T<:Real,N}
    if power(r) < POWER_THRESHOLD || recursion > TRACE_RECURSION_LIMIT
        return nopower
    end
    intsct = closestintersection(assembly, r)
    if intsct === nothing
        return nothing
    else
        surfintsct = point(intsct)
        nml = normal(intsct)
        opticalinterface = interface(intsct)
        λ = wavelength(r)

        if VERSION < v"1.6.0-DEV"
            # TODO REMOVE
            temp = @unionsplit OpticalInterface T opticalinterface processintersection(opticalinterface, surfintsct, nml, r, temperature, pressure, test, recursion == 0)
        else
            temp = processintersection(opticalinterface, surfintsct, nml, r, temperature, pressure, test, recursion == 0)
        end

        if temp === nothing
            return nopower
        else
            raydirection, raypower, raypathlength = temp
            if trackrays !== nothing
                # we want the power that is hitting the surface, i.e. power on incident ray, but the path length at this intersection
                # i.e. with the path length for this intersection added and power modulated by absorption only
                push!(trackrays, LensTrace(OpticalRay(ray(r), raypower, λ, opl = raypathlength, nhits = nhits(r), sourcenum = sourcenum(r), sourcepower = sourcepower(r)), intsct))
            end
            offsetray = OpticalRay(surfintsct + RAY_OFFSET * raydirection, raydirection, raypower, λ, opl = raypathlength, nhits = nhits(r) + 1, sourcenum = sourcenum(r), sourcepower = sourcepower(r))
            res = trace(assembly, offsetray, temperature, pressure, trackrays = trackrays, test = test, recursion = recursion + 1)
            if res === nothing
                return LensTrace(OpticalRay(surfintsct, raydirection, raypower, λ, opl = raypathlength, nhits = nhits(r), sourcenum = sourcenum(r), sourcepower = sourcepower(r)), intsct)
            else
                return res
            end
        end
    end
end

# """approximate focal length of simple double convex lens. For plano convex set one of the radii to Inf"""
focallength(index::Float64, rf::Float64, rb::Float64) = 1 / ((index - 1) * (1 / rf - 1 / rb))