abstract type Trace end
"""
    FactTrace

See [`Trace`](@ref).
"""
struct FactTrace{FT,T,S,S2,R} <: Trace
    F::FT
    t0::T
    x0::S
    θ0::S2
    events::R
end

"""
    FactTrace

See [`Trace`](@ref).
"""
struct PDMPTrace{FT,T,S,S2,S3,R} <: Trace
    F::FT
    t0::T
    x0::S
    θ0::S2
    f::S3
    events::R
end

"""
    Trace(t0::T, x0, θ0, F::Union{ZigZag,FactBoomerang})

Trace object for exact trajectory of pdmp samplers. Returns an iterable `FactTrace` object.
Note that iteration iterates pairs `t => x` where the vector `x` is modified
inplace, so copies have to be made if the `x` is to be saved.
`collect` applied to a trace object automatically copies `x`.
[`discretize`](@ref) returns a discretized version.
"""
Trace(t0::T, x0, θ0, F::Union{ZigZag,FactBoomerang,JointFlow}) where {T} = FactTrace(F, t0, x0, θ0, Tuple{T,Int,eltype(x0),eltype(θ0)}[])
Trace(t0::T, x0::U, θ0::U2, F::Union{BouncyParticle,Boomerang}) where {T, U, U2} = PDMPTrace(F, t0, x0, θ0, ones(Bool, length(x0)), Tuple{T,U,U2,Nothing}[])
Trace(t0::T, x0::U, θ0::U2, f::U3, F::Union{BouncyParticle,Boomerang}) where {T, U, U2, U3} = PDMPTrace(F, t0, x0, θ0, f, Tuple{T,U,U2,U3}[])

Base.length(FT::Trace) = 1 + length(FT.events)

function Base.iterate(FT::FactTrace)
    t, x, θ = FT.t0, copy(FT.x0), copy(FT.θ0)
    t => x, (t, x, θ, 1)
end

function Base.iterate(FT::PDMPTrace)
    t, x, θ = FT.t0, FT.x0, FT.θ0
    t => x, (t, x, θ, 1)
end

function Base.iterate(FT::FactTrace, (t, x, θ, k))
    k > length(FT.events) && return nothing
    k == length(FT.events) && return t => x, (t, x, θ, k + 1)
    t2, i, xi, θi = FT.events[k]
    t, x, θ = move_forward!(t2 - t, t, x, θ, FT.F)
    t = t2
    x[i] = xi
    θ[i] = θi
    return t => x, (t, x, θ, k + 1)
end

function Base.iterate(FT::PDMPTrace, (t, x, θ, k))
    k > length(FT.events) && return nothing
    k == length(FT.events) && return t => x, (t, x, θ, k + 1)
    t, x, θ, _ = FT.events[k]
    return t => x, (t, x, θ, k + 1)
end

Base.push!(FT::Trace, ev) = push!(FT.events, ev)
Base.collect(FT::FactTrace) = collect(t=>copy(x) for (t, x) in FT)
Base.collect(FT::PDMPTrace) = collect(t=>x for (t, x) in FT)

struct Discretize{T, S}
    FT::T
    dt::S
end
Base.IteratorSize(::Discretize) = Iterators.SizeUnknown()

"""
    discretize(trace::Trace, dt)

Discretize `trace` with step-size dt. Returns iterable object
iterating pairs `t => x`.

Iteration changes the vector `x` inplace,
`collect` creates necessary copies.
"""
discretize(FT, dt) = Discretize(FT, dt)

function Base.iterate(D::Discretize{<:FactTrace})
    FT = D.FT
    t, x, θ = FT.t0, copy(FT.x0), copy(FT.θ0)
    t => x, (t, x, θ, 1)
end

function Base.iterate(D::Discretize{<:PDMPTrace})
    FT = D.FT
    t, x, θ, f = FT.t0, copy(FT.x0), copy(FT.θ0), copy(FT.f)
    t => x, (t, x, θ, f, 1)
end

function Base.iterate(D::Discretize{<:FactTrace{T}}, (t, x, θ, k)) where {T<:Union{FactBoomerang,ZigZag,JointFlow}}
    dt = D.dt
    FT = D.FT
    while true
        k > length(FT.events) && return nothing
        ti, i, xi, θi = FT.events[k]
        if t + dt < ti
            t, x, θ = move_forward!(dt, t, x, θ, FT.F)
            return t => x, (t, x, θ, k)
        else # move not more than to ti to change direction
            Δt = ti - t
            dt = dt - Δt
            t, x, θ = move_forward!(Δt, t, x, θ, FT.F)
            t = ti
            x[i] = xi
            θ[i] = θi
            k = k + 1
        end
    end
end

function Base.iterate(D::Discretize{<:PDMPTrace{T}}, (t, x, θ, f, k)) where {T<:Union{Boomerang,BouncyParticle}}
    dt = D.dt
    FT = D.FT
    while true
        k > length(FT.events) && return nothing
        tn = first(FT.events[k])
        if t + dt < tn
            t, x, θ = smove_forward!(dt, t, x, θ, f, FT.F)
            return t => x, (t, x, θ, f, k)
        else # move not more than to ti to change direction
            Δt = tn - t
            dt = dt - Δt
            t, x, θ, f_ = FT.events[k]
            if !isnothing(f_)
                f = f_
            end
            k = k + 1
        end
    end
end

function Base.collect(D::Discretize{<:FactTrace})
    collect(t=>copy(x) for (t, x) in D)
end
function Base.collect(D::Discretize{<:PDMPTrace})
    collect(t=>copy(x) for (t, x) in D)
end

function Statistics.mean(trace::ZigZagBoomerang.Trace)
    x = copy(trace.x0)
    θ = copy(trace.θ0)
    y = 0*x
    T = trace.events[end][1]
    t2 = trace.t0
    t = fill(t2, length(x))
    k = 1
    scale = 1/(2T)
    while k <= length(trace.events)
        t2, i, xi, θi = trace.events[k]
        k += 1
        y[i] += (x[i]+xi)*(t2-t[i])*scale
        t[i] = t2
        x[i] = xi
        θ[i] = θi
    end
    y
end

function cummean(trace::FactTrace)
    x = copy(trace.x0)
    θ = copy(trace.θ0)
    y = 0*x
    T = trace.events[end][1]
    t2 = trace.t0
    ys = [(t=[t2], y=[xi]) for xi in x]

    t = fill(t2, length(x))
    k = 1
    while k <= length(trace.events)
        t2, i, xi, θi = trace.events[k]
        k += 1
        y[i] += (x[i]+xi)*(t2-t[i])
        t[i] = t2
        x[i] = xi
        θ[i] = θi
        push!(ys[i].t, t[i])
        push!(ys[i].y, y[i]/(2t[i]))
    end
    ys
end
export cummean

function Statistics.mean(trace::PDMPTrace)
    x = copy(trace.x0)
    θ = copy(trace.θ0)
    y = 0*x
    ys = fill(y, 0)
    T = trace.events[end][1]
    t2 = trace.t0
    t = t2
    k = 1
    while k <= length(trace.events)
        t2, x2, _ =trace.events[k]
        k += 1
        y  += (x + x2)*(t2-t)
        t = t2
        x  = x2
    end
    y/T
end

function cummean(trace::PDMPTrace)
    x = copy(trace.x0)
    θ = copy(trace.θ0)
    y = 0*x
    ys = fill(y, 0)
    T = trace.events[end][1]
    t2 = trace.t0
    t = t2
    k = 1
    while k <= length(trace.events)
        t2, x2, _ =trace.events[k]
        k += 1
        y  += (x + x2)*(t2-t)
        t = t2
        x  = x2
        push!(ys, y/(2t))
    end
    ys
end

"""
    subtrace(tr, J)

Compute the trace of a subvector `x[J]`,
returns a trace object.
"""
function subtrace(tr, J)
    @assert issorted(J)
    F = tr.F
    t0 = tr.t0
    x0 = tr.x0[J]
    θ0 = tr.θ0[J]

    str = Trace(t0, x0, θ0, F)
    events = str.events
    for ev in tr.events
        r = searchsorted(J, ev[2])
        isempty(r) && continue
        push!(events, (ev[1], r[1], ev[3:end]...))
    end
    str
end


abstract type FEFunction <: CellField end

"""
"""
function get_free_dof_values(f::FEFunction)
  @abstractmethod
end

"""
"""
function get_cell_dof_values(f::FEFunction)
  @abstractmethod
end

function get_cell_dof_values(f::FEFunction,ttrian::Triangulation)
  get_cell_fe_data(get_cell_dof_values,f,ttrian)
end

function get_cell_fe_data(fun,f,ttrian)
  sface_to_data = fun(f)
  strian = get_triangulation(f)
  if strian === ttrian
    return sface_to_data
  end
  @assert is_change_possible(strian,ttrian)
  D = num_cell_dims(strian)
  sglue = get_glue(strian,Val(D))
  tglue = get_glue(ttrian,Val(D))
  get_cell_fe_data(fun,sface_to_data,sglue,tglue)
end

function get_cell_fe_data(fun,sface_to_data,sglue::FaceToFaceGlue,tglue::FaceToFaceGlue)
  mface_to_sface = sglue.mface_to_tface
  tface_to_mface = tglue.tface_to_mface
  mface_to_data = extend(sface_to_data,mface_to_sface)
  tface_to_data = lazy_map(Reindex(mface_to_data),tface_to_mface)
  tface_to_data
end

function get_cell_fe_data(fun,sface_to_data,sglue::FaceToFaceGlue,tglue::SkeletonPair)
  plus = get_cell_fe_data(fun,sface_to_data,sglue,tglue.plus)
  minus = get_cell_fe_data(fun,sface_to_data,sglue,tglue.minus)
  lazy_map(BlockMap(2,[1,2]),plus,minus)
end

"""
"""
function get_fe_space(f::FEFunction)
  @abstractmethod
end

function get_cell_is_dirichlet(f::FEFunction)
  get_cell_is_dirichlet(get_fe_space(f))
end

function get_cell_is_dirichlet(f::FEFunction,ttrian::Triangulation)
  get_cell_is_dirichlet(get_fe_space(f),ttrian)
end

"""
"""
function test_fe_function(f::FEFunction)
  trian = get_triangulation(f)
  free_values = get_free_dof_values(f)
  fe_space = get_fe_space(f)
  @test length(free_values) == num_free_dofs(fe_space)
  cell_values = get_cell_dof_values(f,trian)
  @test length(cell_values) == num_cells(trian)
end

abstract type FESpace <: GridapType end

# Minimal FE interface (used by FEOperator)

"""
"""
function get_free_dof_ids(f::FESpace)
  @abstractmethod
end

"""
"""
num_free_dofs(f::FESpace) = length(get_free_dof_ids(f))

"""
"""
function zero_free_values(f::FESpace)
  V = get_vector_type(f)
  allocate_vector(V,num_free_dofs(f))
end

function get_vector_type(fs::FESpace)
  @abstractmethod
end

"""
"""
function FEFunction(fe::FESpace, free_values)
  @abstractmethod
end

function CellField(fe::FESpace,cell_vals)
  @abstractmethod
end

function get_triangulation(fe::FESpace)
  @abstractmethod
end

# TODO this is quite hacky. Only needed by the zero mean space
function EvaluationFunction(fe::FESpace, free_values)
  FEFunction(fe,free_values)
end

"""
"""
function Base.zero(f::FESpace)
  free_values = zero_free_values(f)
  EvaluationFunction(f,free_values)
end

function get_dof_value_type(f::FESpace)
  get_dof_value_type(get_fe_basis(f),get_fe_dof_basis(f))
end

function get_dof_value_type(cell_shapefuns::CellField,cell_dof_basis::CellDof)
  cell_dof_values = cell_dof_basis(cell_shapefuns)
  eltype(eltype(cell_dof_values))
end

# Extended FEInterface used by Assemblers

"""
"""
function get_cell_dof_ids(f::FESpace)
  @abstractmethod
end

function get_cell_dof_ids(f::FESpace,ttrian::Triangulation)
  get_cell_fe_data(get_cell_dof_ids,f,ttrian)
end

"""
"""
function get_fe_basis(f::FESpace)
  @abstractmethod
end

function get_cell_shapefuns(f::FESpace)
  msg = "get_cell_shapefuns has been removed, use get_fe_basis instead"
  error(msg)
end

"""
"""
function get_fe_dof_basis(f::FESpace)
  @abstractmethod
end

function get_cell_dof_basis(f::FESpace)
  msg = "get_cell_dof_basis has been removed, use get_fe_dof_basis instead"
  error(msg)
end

function get_trial_fe_basis(f::FESpace)
  v = get_fe_basis(f)
  cell_v = get_data(v)
  cell_u = lazy_map(transpose,cell_v)
  similar_fe_basis(v,cell_u,get_triangulation(v),TrialBasis(),DomainStyle(v))
end

function get_cell_shapefuns_trial(f::FESpace)
  msg = "get_cell_shapefuns_trial has been removed, use get_trial_fe_basis instead"
  error(msg)
end

# Basis related

abstract type BasisStyle end
struct TrialBasis <: BasisStyle end
struct TestBasis <: BasisStyle end

abstract type FEBasis <: CellField end
BasisStyle(::T) where T <: FEBasis = BasisStyle(T)
DomainStyle(::T) where T <: FEBasis = DomainStyle(T)

struct SingleFieldFEBasis{BS<:BasisStyle,DS<:DomainStyle} <: FEBasis
  cell_basis::AbstractArray
  trian::Triangulation
  basis_style::BS
  domain_style::DS
  function SingleFieldFEBasis(
    cell_basis::AbstractArray,
    trian::Triangulation,
    basis_style::BasisStyle,
    domain_style::DomainStyle)

    BS = typeof(basis_style)
    DS = typeof(domain_style)
    new{BS,DS}(Fields.MemoArray(cell_basis),trian,basis_style,domain_style)
  end
end

get_data(f::SingleFieldFEBasis) = f.cell_basis
get_triangulation(f::SingleFieldFEBasis) = f.trian
BasisStyle(::Type{SingleFieldFEBasis{BS,DS}}) where {BS,DS} = BS()
DomainStyle(::Type{SingleFieldFEBasis{BS,DS}}) where {BS,DS} = DS()
function CellData.similar_cell_field(f::SingleFieldFEBasis,cell_data,trian,ds::DomainStyle)
  SingleFieldFEBasis(cell_data,trian,BasisStyle(f),ds)
end
function similar_fe_basis(f::SingleFieldFEBasis,cell_data,trian,bs::BasisStyle,ds::DomainStyle)
  SingleFieldFEBasis(cell_data,trian,bs,ds)
end

for fun in (:change_domain_ref_ref,:change_domain_phys_phys)
  @eval begin

    function $fun(
      a::CellData.CellFieldAt{S,<:FEBasis} where S,ttrian::Triangulation,sglue::FaceToFaceGlue,tglue::SkeletonPair)
      a_on_plus_trian = $fun(a.parent,ttrian,sglue,tglue.plus)
      a_on_minus_trian = $fun(a.parent,ttrian,sglue,tglue.minus)
      pair_in = SkeletonPair(get_data(a_on_plus_trian),get_data(a_on_minus_trian))
      pair_out = _fix_cell_basis_dofs_at_skeleton(pair_in,BasisStyle(a.parent))
      if isa(a,CellData.CellFieldAt{:plus})
        return CellData.similar_cell_field(a.parent,pair_out.plus,ttrian,DomainStyle(a.parent))
      elseif isa(a,CellData.CellFieldAt{:minus})
        return CellData.similar_cell_field(a.parent,pair_out.minus,ttrian,DomainStyle(a.parent))
      else
        @unreachable
      end
    end

  end
end

function _fix_cell_basis_dofs_at_skeleton(pair,::TestBasis)
  plus = lazy_map(BlockMap(2,1),pair.plus)
  minus = lazy_map(BlockMap(2,2),pair.minus)
  SkeletonPair(plus,minus)
end

function _fix_cell_basis_dofs_at_skeleton(pair,::TrialBasis)
  plus = lazy_map(BlockMap((1,2),1),pair.plus)
  minus = lazy_map(BlockMap((1,2),2),pair.minus)
  SkeletonPair(plus,minus)
end

# Constraint-related

abstract type ConstraintStyle end
struct Constrained <: ConstraintStyle end
struct UnConstrained <: ConstraintStyle end

ConstraintStyle(::Type{<:FESpace}) = @abstractmethod
ConstraintStyle(::T) where T<:FESpace = ConstraintStyle(T)

has_constraints(::Type{T}) where T <:FESpace = ConstraintStyle(T) == Constrained()
has_constraints(::T) where T <: FESpace = has_constraints(T)

function get_cell_constraints(f::FESpace)
  get_cell_constraints(f,ConstraintStyle(f))
end

function get_cell_constraints(f::FESpace,::UnConstrained)
  cell_ids = get_cell_dof_ids(f)
  cell_axes = lazy_map(axes,cell_ids)
  identity_constraints(cell_axes)
end

function get_cell_constraints(f::FESpace,::Constrained)
  @abstractmethod
end

function get_cell_constraints(f::FESpace,ttrian::Triangulation)
  get_cell_fe_data(get_cell_constraints,f,ttrian)
end

function get_cell_fe_data(fun::typeof(get_cell_constraints),sface_to_data,sglue::FaceToFaceGlue,tglue::SkeletonPair)
  plus = get_cell_fe_data(fun,sface_to_data,sglue,tglue.plus)
  minus = get_cell_fe_data(fun,sface_to_data,sglue,tglue.minus)
  lazy_map(BlockMap((2,2),[(1,1),(2,2)]),plus,minus)
end

function get_cell_isconstrained(f::FESpace)
  get_cell_isconstrained(f,ConstraintStyle(f))
end

function get_cell_isconstrained(f::FESpace,::UnConstrained)
  trian = get_triangulation(f)
  Fill(false,num_cells(trian))
end

function get_cell_isconstrained(f::FESpace,::Constrained)
  @abstractmethod
end

function get_cell_isconstrained(f::FESpace,ttrian::Triangulation)
  get_cell_fe_data(get_cell_isconstrained,f,ttrian)
end

function get_cell_fe_data(
  fun::typeof(get_cell_isconstrained),sface_to_data,sglue::FaceToFaceGlue,tglue::SkeletonPair)
  plus = get_cell_fe_data(fun,sface_to_data,sglue,tglue.plus)
  minus = get_cell_fe_data(fun,sface_to_data,sglue,tglue.minus)
  lazy_map((l,r)-> l||r,plus,minus)
end

function attach_constraints_rows(f::FESpace,cellarr,ttrian::Triangulation)
  attach_constraints_rows(f,cellarr,ttrian,ConstraintStyle(f))
end

function attach_constraints_rows(f::FESpace,cellarr,ttrian,::UnConstrained)
  cellarr
end

function attach_constraints_rows(f::FESpace,cellarr,ttrian,::Constrained)
  cellconstr = get_cell_constraints(f,ttrian)
  cellmask = get_cell_isconstrained(f,ttrian)
  attach_constraints_rows(cellarr,cellconstr,cellmask)
end

function attach_constraints_cols(f::FESpace,cellarr,ttrian::Triangulation)
  attach_constraints_cols(f,cellarr,ttrian,ConstraintStyle(f))
end

function attach_constraints_cols(f::FESpace,cellarr,ttrian,::UnConstrained)
  cellarr
end

function attach_constraints_cols(f::FESpace,cellarr,ttrian,::Constrained)
  cellconstr = get_cell_constraints(f,ttrian)
  cellmask = get_cell_isconstrained(f,ttrian)
  attach_constraints_cols(cellarr,cellconstr,cellmask)
end

function get_cell_is_dirichlet(f::FESpace)
  trian = get_triangulation(f)
  Fill(true,num_cells(trian))
end

function get_cell_is_dirichlet(f::FESpace,ttrian::Triangulation)
  get_cell_fe_data(get_cell_is_dirichlet,f,ttrian)
end

function get_cell_fe_data(
  fun::typeof(get_cell_is_dirichlet),sface_to_data,sglue::FaceToFaceGlue,tglue::SkeletonPair)
  plus = get_cell_fe_data(fun,sface_to_data,sglue,tglue.plus)
  minus = get_cell_fe_data(fun,sface_to_data,sglue,tglue.minus)
  lazy_map((l,r)-> l||r,plus,minus)
end

"""
"""
function test_fe_space(f::FESpace)
  trian = get_triangulation(f)
  @test isa(trian,Triangulation)
  free_values = zero_free_values(f)
  @test length(free_values) == num_free_dofs(f)
  V = get_vector_type(f)
  @test typeof(free_values) == V
  fe_function = FEFunction(f,free_values)
  test_fe_function(fe_function)
  fe_basis = get_fe_basis(f)
  @test isa(has_constraints(f),Bool)
  @test isa(has_constraints(typeof(f)),Bool)
  @test length(get_cell_dof_ids(f,trian)) == num_cells(fe_basis)
  @test length(get_cell_constraints(f,trian)) == num_cells(fe_basis)
  @test length(get_cell_isconstrained(f,trian)) == num_cells(fe_basis)
  @test CellField(f,get_cell_dof_ids(f,trian)) != nothing
end

function test_fe_space(f::FESpace,cell_matvec,cell_mat,cell_vec,trian)
  test_fe_space(f)

  cm = attach_constraints_cols(f,cell_mat,trian)
  if ! has_constraints(f)
    @test cm === cell_mat
  end
  cm = attach_constraints_rows(f,cell_mat,trian)
  if ! has_constraints(f)
    @test cm === cell_mat
  end

  cv = attach_constraints_rows(f,cell_vec,trian)
  if ! has_constraints(f)
    @test cv === cell_vec
  end

  cmv = attach_constraints_cols(f,cell_matvec,trian)
  if ! has_constraints(f)
    @test cmv === cell_matvec
  end
  cmv = attach_constraints_rows(f,cell_matvec,trian)
  if ! has_constraints(f)
    @test cmv === cell_matvec
  end

end


## Convert_a functions for the mdp
# Converts from vector to an action
function POMDPs.convert_a(::Type{Vector{Disturbance}}, avec::AbstractArray{Float64}, mdp::AdversarialDrivingMDP)
    a = Disturbance[]
    index = 1
    for agent in adversaries(mdp)
        d = agent.vec_to_disturbance(avec[index : index + agent.disturbance_dim - 1])
        index += agent.disturbance_dim
        push!(a, d)
    end
    a
end

# Convert from action to a vector
function POMDPs.convert_a(::Type{AbstractArray}, a::Vector{Disturbance}, mdp::AdversarialDrivingMDP)
    vec = Vector{Float64}(undef, sum([agent.disturbance_dim for agent in adversaries(mdp)]))
    index = 1
    for i = 1:length(adversaries(mdp))
        agent = adversaries(mdp)[i]
        vec[index : index + agent.disturbance_dim - 1] .= agent.disturbance_to_vec(a[i])
        index += agent.disturbance_dim
    end
    vec
end

## PedestrianControl Disturbances
const PEDESTRIAN_DISTURBANCE_DIM = 5

# Converts PedestrianControl disturbance to a vector
function PedestrianControl_to_vec(pc::PedestrianControl)
    Float64[pc.da..., pc.noise.pos..., pc.noise.vel]
end

# Converts from a vector to a PedestrianConrol Disturbance
function vec_to_PedestrianControl(arr::AbstractArray)
    @assert length(arr) == PEDESTRIAN_DISTURBANCE_DIM
    da = VecE2(arr[1], arr[2])
    noise_pos = VecE2(arr[3], arr[4])
    noise_v = arr[5]
    PedestrianControl(VecE2(0., 0.), da, Noise(noise_pos, noise_v))
end

## BlinkerVehicle Disturbances
const BLINKERVEHICLE_DISTURBANCE_DIM = 5

# Converts PedestrianControl disturbance to a vector
function BlinkerVehicleControl_to_vec(bv::BlinkerVehicleControl)
    g = Float64(bv.toggle_goal)
    b = Float64(bv.toggle_blinker)
    Float64[bv.da, g, b, bv.noise.pos[1], bv.noise.vel]
end

# Converts from a vector to a PedestrianConrol Disturbance
function vec_to_BlinkerVehicleControl(arr::AbstractArray)
    @assert length(arr) == BLINKERVEHICLE_DISTURBANCE_DIM
    da = arr[1]
    toggle_goal = arr[2] > 0
    toggle_blinker = arr[3] > 0
    noise_pos = VecE2(arr[4], 0.)
    noise_v = arr[5]
    BlinkerVehicleControl(0., da, toggle_goal, toggle_blinker, Noise(noise_pos, noise_v))
end

## Continuous action spaces
get_actions(m::Vector{Sampleable}) = error("Error! `actions` called on continuous action space")
get_actionindex(m::Vector{Sampleable}, a::Vector{Disturbance}) = error("Error! `actionindex` called on continuous action space")
function Distributions.logpdf(m::Vector{Sampleable}, a::Vector{Disturbance}, mdp::AdversarialDrivingMDP)
    avec = convert_a(AbstractArray, a, mdp)
    sum([logpdf(m[i], avec[i]) for i=1:length(avec)])
end
Base.rand(rng::AbstractRNG, m::Vector{Sampleable}, mdp::AdversarialDrivingMDP) = convert_a(Vector{Disturbance}, [rand(rng, d) for d in m], mdp)
Base.rand(rng::AbstractRNG, m::Vector{Sampleable}, convert_fn::Function) = convert_fn([rand(rng, d) for d in m])

## Discrete BlinkerVehicle Actions
struct DiscreteActionModel
    actions::Vector{Vector{Disturbance}}
    actionindex::Dict{Vector{Disturbance}, Int}
    probs::Vector{Float64}
end

function DiscreteActionModel(actions::Vector{Vector{Disturbance}}, probs::Vector{Float64})
    aind = Dict{Vector{Disturbance}, Int}(a => i for (a, i) in zip(actions, 1:length(actions)))
    DiscreteActionModel(actions, aind, probs)
end

get_actions(m::DiscreteActionModel) = m.actions
get_actionindex(m::DiscreteActionModel, a::Vector{Disturbance}) = m.actionindex[a]
Distributions.logpdf(m::DiscreteActionModel, a::Vector{Disturbance}, mdp::AdversarialDrivingMDP) = log(m.probs[findfirst(m.actions .== [a])])
Base.rand(rng::AbstractRNG, m::DiscreteActionModel, mdp::AdversarialDrivingMDP) = m.actions[rand(rng, Categorical(m.probs))]

## Combine disturbance_models
# Automatically determine the type of disturbance (continuous or discrete)
combine_disturbance_models(adversaries::Vector{Agent}) = adversaries[1].disturbance_model isa DiscreteActionModel ? combine_discrete(adversaries) : combine_continuous(adversaries)

# Combines continuous disturbance models
function combine_continuous(adversaries::Vector{Agent})
    m = Vector{Sampleable}()
    for a in adversaries
        push!(m, a.disturbance_model...)
    end
    m
end

# Combines discrete disturbance models
function combine_discrete(adversaries::Vector{Agent})
    actions = Vector{Disturbance}[]
    action_probabilities = Float64[]

    # Add the baseline action where all agents do nothing
    base_action = Disturbance[a.disturbance_model.actions[1][1] for a in adversaries]
    aprob = sum([a.disturbance_model.probs[1] for a in adversaries])
    push!(actions, base_action)
    push!(action_probabilities, aprob)

    # Loops through the rest of the actions
    for adv_i=1:length(adversaries)
        adv = adversaries[adv_i]
        for (act, prob) in zip(adv.disturbance_model.actions[2:end], adv.disturbance_model.probs[2:end])
            new_action = deepcopy(base_action)
            new_action[adv_i] = act[1]
            push!(actions, new_action)
            push!(action_probabilities, prob)
        end
    end
    action_probabilities ./= sum(action_probabilities)
    DiscreteActionModel(actions, action_probabilities)
end

## default actions
# Function to setup default actions for the blinker verhicle
get_bv_actions(med_accel = 1.5, large_accel = 3.0, med_prob = 1e-2, large_prob = 1e-3) = DiscreteActionModel(
                      Vector{Vector{Disturbance}}([ [BlinkerVehicleControl(0, 0., false, false, Noise())],
                        [BlinkerVehicleControl(0, -large_accel, false, false, Noise())],
                        [BlinkerVehicleControl(0, -med_accel, false, false, Noise())],
                        [BlinkerVehicleControl(0, med_accel, false, false, Noise())],
                        [BlinkerVehicleControl(0, large_accel, false, false, Noise())],
                        [BlinkerVehicleControl(0, 0., true, false, Noise())], # toggle goal
                        [BlinkerVehicleControl(0, 0., false, true, Noise())] # toggle blinker
                       ]),
                        [1 - (4*large_prob + 2*med_prob), large_prob, med_prob, med_prob, large_prob, large_prob, large_prob]
                        )

function get_rand_bv_actions(rng; med_accel = 1.5, large_accel = 3.0)
    probs = 10 .^ [rand(rng, Uniform(-4, -2)),
                   rand(rng, Uniform(-3, -1)),
                   rand(rng, Uniform(-3, -1)),
                   rand(rng, Uniform(-4, -2)),
                   rand(rng, Uniform(-4, -2)),
                   rand(rng, Uniform(-4, -2))]
    leftover = 1.0 - sum(probs)
    DiscreteActionModel( Vector{Vector{Disturbance}}([ [BlinkerVehicleControl(0, 0., false, false, Noise())],
                            [BlinkerVehicleControl(0, -large_accel, false, false, Noise())],
                            [BlinkerVehicleControl(0, -med_accel, false, false, Noise())],
                            [BlinkerVehicleControl(0, med_accel, false, false, Noise())],
                            [BlinkerVehicleControl(0, large_accel, false, false, Noise())],
                            [BlinkerVehicleControl(0, 0., true, false, Noise())], # toggle goal
                            [BlinkerVehicleControl(0, 0., false, true, Noise())] # toggle blinker
                           ]),
                            [leftover, probs...]
                            )
end

# TODO Add default distributions as needed


### Working with cooler files, uses cool https://github.com/open2c/cooler

using PyCall
using DataFrames
using OffsetArrays, CoordinateTransformations, Rotations
using ImageTransformations

loadcooler() = pyimport("cooler")
loadcoolerfile(coolerfile; cooler=loadcooler()) = cooler.Cooler(coolerfile)

### load bin dirs
coolercoords(chrom, loc) = (replace(chrom, "chr" => ""), first(loc), last(loc))
function bindf(chrom, loc, cool; ext=0.5, saveload=true, projdir=getprojectdir())

    dffile =  joinpath(projdir, "results", "cooler", string("isletcooler_bindf_", chrom, "_", first(loc), "_", last(loc), "", ".tsv.gz"))
    if isnothing(cool) || !saveload
        if !isfile(dffile)
            error("Error: cannot load precalculated contact matrix bins $dffile, please load cooler file")
        end
        return CSV.read(dffile, DataFrame)
    end

    dx = Int(round((last(loc) - first(loc))*ext))
    xloc = (first(loc) - dx):(last(loc) + dx)

    bdf = cool.bins().fetch(coolercoords(chrom, xloc))
    cols = Symbol.(bdf.columns)
    data = [[v for v in bdf[c]] for c in cols]
    df = DataFrame(data, cols)
    df.binmid = div.(df.start .+ df.end, 2)
    df

    if saveload
        mkpath(dirname(dffile))
        CSV.write(dffile, df, compress=true)
    end
    df
end

## cooler matrix
function coolmatrix(chrom, loc, cool; balance="KR", ext=0.5, saveload=true, projdir=getprojectdir())

    matfile =  joinpath(projdir, "results", "cooler", string("isletcooler_", chrom, "_", first(loc), "_", last(loc), "", ".tsv.gz"))
    if isnothing(cool) || !saveload
        if !isfile(matfile)
            error("Error: cannot load precalculated contact matrix file $matfile, please load cooler file")
        end
        return Matrix{Float64}(CSV.read(matfile, DataFrame)[!, 2:end])
    end

    dx = Int(round((last(loc) - first(loc))*ext))
    xloc = (first(loc) - dx):(last(loc) + dx)
    CM = cool.matrix(balance=balance).fetch(coolercoords(chrom, xloc))

    if saveload
        mkpath(dirname(matfile))
        CSV.write(matfile, DataFrame(CM, :auto), compress=true)
    end
    CM
end

### rotate matrix to be 2d
function rotcooler(xp, CM ; kwargs...)
    R = warp(CM, recenter(RotMatrix(-π/4), [0, 0]) ; kwargs...)
    n = cld(size(R, 1), 2)
    R = R.parent[(n+1):end, :]
    nxp = range(first(xp), last(xp), length=size(R, 2))
    dx = xp[2] - xp[1]
    nyp = dx*(1:size(R, 1))
    nxp, nyp, R
end

## quantify islet loops
function quantisletloops(hicdata, cool, binsize=5000, ext=2, saveload=true, projdir=getprojectdir())
    datafile = joinpath(projdir, "results", "cooler", string("islet_hicquant.tsv.gz"))
    if isnothing(cool) || !saveload
        if !isfile(datafile)
            error("Error: cannot load precalculated quantification of islet loops $datafile, please load cooler file")
        end
        return CSV.read(datafile, DataFrame)
    end

    dx = binsize*ext
    μ = Float64[]
    @showprogress for row in eachrow(hicdata)
        locA = (row.startA - dx):(row.stopA + dx)
        locB = (row.startB - dx):(row.stopB + dx)
        m = cool.matrix(balance="KR").fetch(coolercoords(row.chrA, locA), coolercoords(row.chrB, locB))
        push!(μ, mean(m))
    end

    hicdata.LoopQuant = μ
    if saveload
        mkpath(dirname(datafile))
        CSV.write(datafile, hicdata, delim='\t', compress=true)
    end
    hicdata
end

### Plot HiC matrix over a region
function coolerpanel(chrom, loc, isletcooler, hicdata, regioncoords;  ext=0.75, plotroi=true)

    bincoords = bindf(chrom, loc, isletcooler, ext=ext)
    CM = coolmatrix(chrom, loc, isletcooler, ext=ext)
    nxp, nyp, R = rotcooler(bincoords.binmid, CM)
    cl=(0, 2.3139)
    ph = heatmap(nxp, nyp, log10.(R .+ 1), ylims=(0, 4.0e+5), colorbar=false, c=:inferno, xticks=false, clims=cl) ## clims hardcoded to compare to HK1 scale
    # cl = zlims()
    # @show cl
    yticks!((0:200:400)*1e+3, string.(0:.2:.4))

    hicloops = @subset(hicdata, :Study .== "IsletHiC", :chrA .== chrom, (:startA .∈ Ref(loc)) .| (:stopB .∈ Ref(loc)) .| (:startB .∈ Ref(loc)) .| (:stopA .∈ Ref(loc))) ## islet loops
    plottriloops!(hicloops, c=:white, lab="Islet Loops", plotlines=true)

    plot!(ylabel="Distance (Mb)")

    plotroi && vline!([mean(last(regioncoords["roi"]))], c=:white, lab="ROI", ls=:dash, title="Islets HiC")
    pg = plotgenemodels(loc, genecoords, fs=6)
    plot!(ylims=(-2, 2), yticks=false, yaxis=false, xticks=true, xlabel=string(chrom, ":", loc[1], "-", loc[2]), bottom_margin=5mm)
    plotroi && vline!([mean(last(hk1coords["roi"]))], c=:black, lab="", ls=:dash)
    pblank = plot([0], [0], zcolor=[NaN], colorbar=true, clims=cl, lab="", framestyle=:none, c=:inferno)
    lt = @layout [[a{0.65h}  ; b] c{0.05w}]

    p = plot(ph, pg, layout=lt, pblank, size=(900, 290), link=:x, xlims=(first(loc), last(loc)), fmt=:png, background_color_legend=:black, foreground_color_legend=:white, legendfontcolor=:white,
    left_margin=5mm, fontfamily="helvetica")

end

### plot triangluar loops over HiC Matrix
function plottriloops!(df; w=15000, np=100, c=:white, ss=1, plotlines=true, lab="")

    ps = Plots.partialcircle(0, 2π, np, w)
    ofs = ss*w/sqrt(2)

    for (k, row) in enumerate(eachrow(df))
        l = (row.startA + row.stopA)/2
        r = (row.startB + row.stopB)/2

        m, s = coordconv(l, r)
        plotlines && plot!([l, m - ofs], [0, s-ofs*sqrt(2)], c=c, lab="")
        plotlines && plot!([m+ofs, r], [s-ofs*sqrt(2), 0], c=c, lab="")
        plot!(first.(ps) .+ m, last.(ps).*sqrt(2) .+ s, c=c, lab=ifelse(k == 1, lab, ""))
    end
    plot!()
end

function coordconv(start, stop)
    m = (start+stop)/2
    s = (stop - start)/sqrt(2)
    m, s
end

iv(a, b) = a:b
function hk1_loop_summary(hicdata, hk1coords ; trf = x -> log10(x + 1))

    chrom, loc = hk1coords["roi"]
    ### hk1 loop

    hk1loops = @subset(hicdata, :chrA .== chrom, :chrB .== chrom, .!isempty.(intersect.(iv.(:startA, :stopB), Ref(loc))),  ((:stopB .- :startA) .< 1e+6))
    hk1loops.roidist = abs.(first(loc) .- hk1loops.startA) .+ abs.(last(loc) .- hk1loops.startB)
    sort!(hk1loops, :roidist) # take quantification of loop that tightest around the ROI
    lq = trf(first(hk1loops.LoopQuant))

    bins = range(0, 2.25, length=80)
    q = trf.(hicdata.LoopQuant[.!isnan.(hicdata.LoopQuant)])

    ph = stephist(q, bins = bins, c=:steelblue, fill=0, fillalpha=0.2, lab="", ylabel="Frequency", title="Histogram of log Contact Frequencies\nAll Loops Human Islets")
    vline!([lq], c=:black, lab="HK1 Loop", leg=:topleft, ls=:dash)
    ## ecdf
    ec = ecdf(q)

    pe = plot(bins, ec(bins), c=:steelblue, lab="", ylabel="Cumlative Proprtion", title="Empirical cdf log Contact Frequencies\nAll Loops Human Islets")
    yl = ylims()
    xl = xlims()

    plot!([lq, lq], [first(yl), ec(lq)], c=:black, lab="", ls=:dash)
    plot!([first(xl), lq], [ec(lq), ec(lq)], c=:black, lab=string("HK1 Loop: ", round(ec(lq), digits=2)), ls=:dash, leg=:left)
    plot!(xlims=xl, ylims=yl)
    plot(ph, pe, size=(790, 250), xlabel="log10 Loop Contact frequencies", bottom_margin=5mm, fontfamily="helvetica", left_margin=5mm, titlefont=font(12, "helvetica"), top_margin=5mm)

end


struct ODEAdjointSensitivityFunction{rateType,uType,uType2,UF,PF,G,JC,GC,A,DG,TJ,PJT,PJC,CP,SType,INT,CV} <: SensitivityFunction
  uf::UF
  pf::PF
  g::G
  J::TJ
  pJ::PJT
  dg_val::uType
  jac_config::JC
  g_grad_config::GC
  paramjac_config::PJC
  alg::A
  f_cache::rateType
  discrete::Bool
  y::uType2
  sol::SType
  dg::DG
  checkpoints::CP
  integrator::INT
  colorvec::CV
end

@noinline function ODEAdjointSensitivityFunction(g,u0,p,alg,discrete,sol,dg,checkpoints,tspan,colorvec)
  numparams = length(p)
  numindvar = length(u0)
  # if there is an analytical Jacobian provided, we are not going to do automatic `jac*vec`
  f = sol.prob.f
  isautojacvec = DiffEqBase.has_jac(f) ? false : get_jacvec(alg)
  J = isautojacvec ? nothing : similar(u0, numindvar, numindvar)

  if !discrete
    if dg != nothing || isautojacvec
      pg = nothing
      pg_config = nothing
    else
      pg = UGradientWrapper(g,tspan[2],p)
      pg_config = build_grad_config(alg,pg,u0,p)
    end
  else
    pg = nothing
    pg_config = nothing
  end

  if DiffEqBase.has_jac(f) || isautojacvec
    jac_config = nothing
    uf = nothing
  else
    uf = DiffEqDiffTools.UJacobianWrapper(f,tspan[2],p)
    jac_config = build_jac_config(alg,uf,u0)
  end

  y = copy(sol(tspan[1])) # TODO: Has to start at interpolation value!
  paramjac_config = nothing
  pf = nothing
  if !isquad(alg)
    if DiffEqBase.has_paramjac(f) || isautojacvec
      paramjac_config = nothing
    else
      pf = DiffEqDiffTools.ParamJacobianWrapper(f,tspan[1],y)
      paramjac_config = build_param_jac_config(alg,pf,y,p)
    end
  end

  pJ = if !isquad(alg)
    isautojacvec ? nothing : similar(sol.prob.u0, numindvar, numparams)
  else
    nothing
  end
  integrator = if ischeckpointing(alg)
    integ = init(sol.prob, sol.alg, save_on=false)
    integ.u = y
    integ
  else
    nothing
  end
  dg_val = similar(u0, numindvar) # number of funcs size
  f_cache = deepcopy(u0)

  return ODEAdjointSensitivityFunction(uf,pf,pg,J,pJ,dg_val,
                               jac_config,pg_config,paramjac_config,
                               alg,f_cache,
                               discrete,y,sol,dg,checkpoints,integrator,colorvec)
end

# u = λ'
function (S::ODEAdjointSensitivityFunction)(du,u,p,t)
  @unpack y, sol, J, uf, alg, f_cache, jac_config, discrete, dg, dg_val, g, g_grad_config = S
  idx = length(y)
  f = sol.prob.f
  isautojacvec = DiffEqBase.has_jac(f) ? false : get_jacvec(alg)

  if isbcksol(alg)
    λ     = @view u[1:idx]
    dλ    = @view du[1:idx]
    grad  = @view u[idx+1:end-idx]
    dgrad = @view du[idx+1:end-idx]
    _y    = @view u[end-idx+1:end]
    dy    = @view du[end-idx+1:end]
    copyto!(vec(y), _y)
    isautojacvec || f(dy, _y, p, t)
  else
    if ischeckpointing(alg)
      @unpack integrator, checkpoints = S
      # assuming that in the forward direction `t0` < `t1`, and the
      # `checkpoints` vector is sorted with respect to the forward direction
      tidx = findlast(x->x <= t, checkpoints)
      t0 = checkpoints[tidx]
      dt = t-t0
      if abs(dt) > integrator.opts.dtmin
        sol(integrator.u, t0)
        copyto!(integrator.uprev, integrator.u)
        integrator.t = t0
        # set `iter` to some arbitrary integer so that there won't be max maxiters error
        integrator.iter=100
        u_modified!(integrator, true)
        step!(integrator, dt, true)
        # `integrator.u` is aliased to `y`
      else
        sol(y,t)
      end
    else
      sol(y,t)
    end
    if isquad(alg)
      λ     = u
      dλ    = du
    else
      λ     = @view u[1:idx]
      dλ    = @view du[1:idx]
      grad  = @view u[idx+1:end]
      dgrad = @view du[idx+1:end]
    end
  end

  if !isautojacvec
    if DiffEqBase.has_jac(f)
      f.jac(J,y,p,t) # Calculate the Jacobian into J
    else
      uf.t = t
      jacobian!(J, uf, y, f_cache, alg, jac_config)
    end
    mul!(dλ',λ',J)
  elseif isquad(alg)
    _dy, back = Tracker.forward(y) do u
      if DiffEqBase.isinplace(sol.prob)
        out_ = map(zero, u)
        f(out_, u, p, t)
        Tracker.collect(out_)
      else
        vec(f(u, p, t))
      end
    end
    dλ[:] = Tracker.data(back(λ)[1])
    isbcksol(alg) && (dy[:] = vec(Tracker.data(_dy)))
  else
    _dy, back = Tracker.forward(y, sol.prob.p) do u, p
      if DiffEqBase.isinplace(sol.prob)
        out_ = map(zero, u)
        f(out_, u, p, t)
        Tracker.collect(out_)
      else
        vec(f(u, p, t))
      end
    end
    dλ[:], dgrad[:] = Tracker.data.(back(λ))
    isbcksol(alg) && (dy[:] = vec(Tracker.data(_dy)))
  end

  dλ .*= -one(eltype(λ))

  if !discrete
    if dg != nothing
      dg(dg_val,y,p,t)
    else
      g.t = t
      gradient!(dg_val, g, y, alg, g_grad_config)
    end
    dλ .+= dg_val
  end

  if !isquad(alg) && !isautojacvec
    @unpack pJ, pf, paramjac_config = S
    if DiffEqBase.has_paramjac(f)
      f.paramjac(pJ,y,sol.prob.p,t) # Calculate the parameter Jacobian into pJ
    else
      jacobian!(pJ, pf, sol.prob.p, f_cache, alg, paramjac_config)
    end
    mul!(dgrad',λ',pJ)
  end
  nothing
end

# g is either g(t,u,p) or discrete g(t,u,i)
@noinline function ODEAdjointProblem(sol,g,t=nothing,dg=nothing,
                           alg=SensitivityAlg();
                           checkpoints=sol.t,
                           callback=CallbackSet())
  f = sol.prob.f
  tspan = (sol.prob.tspan[2],sol.prob.tspan[1])
  discrete = t != nothing

  p = sol.prob.p
  p === DiffEqBase.NullParameters() && error("Your model does not have parameters, and thus it is impossible to calculate the derivative of the solution with respect to the parameters. Your model must have parameters to use parameter sensitivity calculations!")

  u0 = zero(sol.prob.u0)

  len = isquad(alg) ? length(u0) : length(u0)+length(p)
  λ = similar(u0, len)
  sense = ODEAdjointSensitivityFunction(g,u0,
                                        p,alg,discrete,
                                        sol,dg,checkpoints,tspan,f.colorvec)

  init_cb = t !== nothing && sol.prob.tspan[2] == t[end]
  cb = generate_callbacks(sense, g, λ, t, callback, init_cb)
  z0 = isbcksol(alg) ? [vec(zero(λ)); vec(sense.y)] : vec(zero(λ))
  ODEProblem(sense,z0,tspan,p,callback=cb)
end

function generate_callbacks(sensefun, g, λ, t, callback, init_cb)
  if sensefun.discrete
    @unpack alg, y, sol = sensefun
    prob = sol.prob
    cur_time = Ref(length(t))
    function time_choice(integrator)
      cur_time[] > 0 ? t[cur_time[]] : nothing
    end
    affect! = let isq = isquad(alg), λ=λ, t=t, y=y, cur_time=cur_time, idx=length(prob.u0)
      function (integrator)
        p, u = integrator.p, integrator.u
        λ  = isq ? λ : @view(λ[1:idx])
        g(λ,y,p,t[cur_time[]],cur_time[])
        if isq
          u .+= λ
        else
          u = @view u[1:idx]
          u .= λ .+ @view integrator.u[1:idx]
        end
        u_modified!(integrator,true)
        cur_time[] -= 1
      end
    end
    cb = IterativeCallback(time_choice,affect!,eltype(prob.tspan);initial_affect=init_cb)

    _cb = CallbackSet(cb,callback)
  else
    _cb = callback
  end
  return _cb
end

struct AdjointSensitivityIntegrand{pType,uType,rateType,S,AS,PF,PJC,A,PJT}
  sol::S
  adj_sol::AS
  p::pType
  y::uType
  λ::uType
  pf::PF
  f_cache::rateType
  pJ::PJT
  paramjac_config::PJC
  alg::A
end

function AdjointSensitivityIntegrand(sol,adj_sol,alg=SensitivityAlg())
  prob = sol.prob
  @unpack f, p, tspan, u0 = prob
  y = similar(sol.prob.u0)
  λ = similar(adj_sol.prob.u0)
  # we need to alias `y`
  pf = DiffEqDiffTools.ParamJacobianWrapper(f,tspan[1],y)
  f_cache = similar(y)
  isautojacvec = DiffEqBase.has_paramjac(f) ? false : get_jacvec(alg)
  pJ = isautojacvec ? nothing : similar(u0,length(u0),length(p))

  if DiffEqBase.has_paramjac(f) || isautojacvec
    paramjac_config = nothing
  else
    paramjac_config = build_param_jac_config(alg,pf,y,p)
  end
  AdjointSensitivityIntegrand(sol,adj_sol,p,y,λ,pf,f_cache,pJ,paramjac_config,alg)
end

function (S::AdjointSensitivityIntegrand)(out,t)
  @unpack y, λ, pJ, pf, p, f_cache, paramjac_config, alg, sol, adj_sol = S
  f = sol.prob.f
  sol(y,t)
  adj_sol(λ,t)
  λ .*= -one(eltype(λ))
  isautojacvec = DiffEqBase.has_paramjac(f) ? false : get_jacvec(alg)
  # y is aliased
  pf.t = t

  if !isautojacvec
    if DiffEqBase.has_paramjac(f)
      f.paramjac(pJ,y,p,t) # Calculate the parameter Jacobian into pJ
    else
      jacobian!(pJ, pf, p, f_cache, alg, paramjac_config)
    end
    mul!(out',λ',pJ)
  else
    _, back = Tracker.forward(y, p) do u, p
      if DiffEqBase.isinplace(sol.prob)
        out_ = map(zero, u)
        f(out_, u, p, t)
        Tracker.collect(out_)
      else
        vec(f(u, p, t))
      end
    end
    out[:] = vec(Tracker.data(back(λ)[2]))
  end
  out'
end

function (S::AdjointSensitivityIntegrand)(t)
  out = similar(S.p)
  S(out,t)
end

function adjoint_sensitivities_u0(sol,alg,g,t=nothing,dg=nothing;
                                  abstol=1e-6,reltol=1e-3,
                                  iabstol=abstol, ireltol=reltol,sensealg=SensitivityAlg(checkpointing=!sol.dense,quad=false),
                                  checkpoints=sol.t,
                                  kwargs...)
  isquad(sensealg) && error("Can't get sensitivities of u0 with quadrature.")
  adj_prob = ODEAdjointProblem(sol,g,t,dg,sensealg,checkpoints=checkpoints)
  adj_sol = solve(adj_prob,alg;abstol=abstol,reltol=reltol,kwargs...,save_everystep=false,save_start=false,saveat=eltype(sol[1])[])

  -adj_sol[end][1:length(sol.prob.u0)],
    adj_sol[end][(1:length(sol.prob.p)) .+ length(sol.prob.u0)]'
end

function adjoint_sensitivities(sol,alg,g,t=nothing,dg=nothing;
                               abstol=1e-6,reltol=1e-3,
                               iabstol=abstol, ireltol=reltol,sensealg=SensitivityAlg(checkpointing=!sol.dense),
                               checkpoints=sol.t,
                               kwargs...)
  adj_prob = ODEAdjointProblem(sol,g,t,dg,sensealg,checkpoints=checkpoints)
  isq = isquad(sensealg)
  adj_sol = solve(adj_prob,alg;abstol=abstol,reltol=reltol,
                               save_everystep=isq,save_start=isq,kwargs...)
  !isq && return adj_sol[end][(1:length(sol.prob.p)) .+ length(sol.prob.u0)]'
  integrand = AdjointSensitivityIntegrand(sol,adj_sol,sensealg)

  if t === nothing
    res,err = quadgk(integrand,sol.prob.tspan[1],sol.prob.tspan[2],
                   atol=iabstol,rtol=ireltol)
  else
    res = zero(integrand.p)'
    for i in 1:length(t)-1
      res .+= quadgk(integrand,t[i],t[i+1],
                     atol=iabstol,rtol=ireltol)[1]
    end
    if t[1] != sol.prob.tspan[1]
      res .+= quadgk(integrand,sol.prob.tspan[1],t[1],
                     atol=iabstol,rtol=ireltol)[1]
    end
  end
  res
end


"""
    space(::SiteType"Electron";
          conserve_qns = false,
          conserve_sz = conserve_qns,
          conserve_nf = conserve_qns,
          conserve_nfparity = conserve_qns,
          qnname_sz = "Sz",
          qnname_nf = "Nf",
          qnname_nfparity = "NfParity")

Create the Hilbert space for a site of type "Electron".

Optionally specify the conserved symmetries and their quantum number labels.
"""
function ITensors.space(
  ::SiteType"Electron";
  conserve_qns=false,
  conserve_sz=conserve_qns,
  conserve_nf=conserve_qns,
  conserve_nfparity=conserve_qns,
  qnname_sz="Sz",
  qnname_nf="Nf",
  qnname_nfparity="NfParity",
  # Deprecated
  conserve_parity=nothing,
)
  if !isnothing(conserve_parity)
    conserve_nfparity = conserve_parity
  end
  if conserve_sz && conserve_nf
    return [
      QN((qnname_nf, 0, -1), (qnname_sz, 0)) => 1
      QN((qnname_nf, 1, -1), (qnname_sz, +1)) => 1
      QN((qnname_nf, 1, -1), (qnname_sz, -1)) => 1
      QN((qnname_nf, 2, -1), (qnname_sz, 0)) => 1
    ]
  elseif conserve_nf
    return [
      QN(qnname_nf, 0, -1) => 1
      QN(qnname_nf, 1, -1) => 2
      QN(qnname_nf, 2, -1) => 1
    ]
  elseif conserve_sz
    return [
      QN((qnname_sz, 0), (qnname_nfparity, 0, -2)) => 1
      QN((qnname_sz, +1), (qnname_nfparity, 1, -2)) => 1
      QN((qnname_sz, -1), (qnname_nfparity, 1, -2)) => 1
      QN((qnname_sz, 0), (qnname_nfparity, 0, -2)) => 1
    ]
  elseif conserve_nfparity
    return [
      QN(qnname_nfparity, 0, -2) => 1
      QN(qnname_nfparity, 1, -2) => 2
      QN(qnname_nfparity, 0, -2) => 1
    ]
  end
  return 4
end

ITensors.val(::ValName"Emp", ::SiteType"Electron") = 1
ITensors.val(::ValName"Up", ::SiteType"Electron") = 2
ITensors.val(::ValName"Dn", ::SiteType"Electron") = 3
ITensors.val(::ValName"UpDn", ::SiteType"Electron") = 4
ITensors.val(::ValName"0", st::SiteType"Electron") = val(ValName("Emp"), st)
ITensors.val(::ValName"↑", st::SiteType"Electron") = val(ValName("Up"), st)
ITensors.val(::ValName"↓", st::SiteType"Electron") = val(ValName("Dn"), st)
ITensors.val(::ValName"↑↓", st::SiteType"Electron") = val(ValName("UpDn"), st)

ITensors.state(::StateName"Emp", ::SiteType"Electron") = [1.0, 0, 0, 0]
ITensors.state(::StateName"Up", ::SiteType"Electron") = [0.0, 1, 0, 0]
ITensors.state(::StateName"Dn", ::SiteType"Electron") = [0.0, 0, 1, 0]
ITensors.state(::StateName"UpDn", ::SiteType"Electron") = [0.0, 0, 0, 1]
ITensors.state(::StateName"0", st::SiteType"Electron") = state(StateName("Emp"), st)
ITensors.state(::StateName"↑", st::SiteType"Electron") = state(StateName("Up"), st)
ITensors.state(::StateName"↓", st::SiteType"Electron") = state(StateName("Dn"), st)
ITensors.state(::StateName"↑↓", st::SiteType"Electron") = state(StateName("UpDn"), st)

function ITensors.op!(Op::ITensor, ::OpName"Nup", ::SiteType"Electron", s::Index)
  Op[s' => 2, s => 2] = 1.0
  return Op[s' => 4, s => 4] = 1.0
end

function ITensors.op!(Op::ITensor, ::OpName"Ndn", ::SiteType"Electron", s::Index)
  Op[s' => 3, s => 3] = 1.0
  return Op[s' => 4, s => 4] = 1.0
end

function ITensors.op!(Op::ITensor, ::OpName"Nupdn", ::SiteType"Electron", s::Index)
  return Op[s' => 4, s => 4] = 1.0
end

function ITensors.op!(Op::ITensor, ::OpName"Ntot", ::SiteType"Electron", s::Index)
  Op[s' => 2, s => 2] = 1.0
  Op[s' => 3, s => 3] = 1.0
  return Op[s' => 4, s => 4] = 2.0
end

function ITensors.op!(Op::ITensor, ::OpName"Cup", ::SiteType"Electron", s::Index)
  Op[s' => 1, s => 2] = 1.0
  return Op[s' => 3, s => 4] = 1.0
end

function ITensors.op!(Op::ITensor, ::OpName"Cdagup", ::SiteType"Electron", s::Index)
  Op[s' => 2, s => 1] = 1.0
  return Op[s' => 4, s => 3] = 1.0
end

function ITensors.op!(Op::ITensor, ::OpName"Cdn", ::SiteType"Electron", s::Index)
  Op[s' => 1, s => 3] = 1.0
  return Op[s' => 2, s => 4] = -1.0
end

function ITensors.op!(Op::ITensor, ::OpName"Cdagdn", ::SiteType"Electron", s::Index)
  Op[s' => 3, s => 1] = 1.0
  return Op[s' => 4, s => 2] = -1.0
end

function ITensors.op!(Op::ITensor, ::OpName"Aup", ::SiteType"Electron", s::Index)
  Op[s' => 1, s => 2] = 1.0
  return Op[s' => 3, s => 4] = 1.0
end

function ITensors.op!(Op::ITensor, ::OpName"Adagup", ::SiteType"Electron", s::Index)
  Op[s' => 2, s => 1] = 1.0
  return Op[s' => 4, s => 3] = 1.0
end

function ITensors.op!(Op::ITensor, ::OpName"Adn", ::SiteType"Electron", s::Index)
  Op[s' => 1, s => 3] = 1.0
  return Op[s' => 2, s => 4] = 1.0
end

function ITensors.op!(Op::ITensor, ::OpName"Adagdn", ::SiteType"Electron", s::Index)
  Op[s' => 3, s => 1] = 1.0
  return Op[s' => 4, s => 2] = 1.0
end

function ITensors.op!(Op::ITensor, ::OpName"F", ::SiteType"Electron", s::Index)
  Op[s' => 1, s => 1] = +1.0
  Op[s' => 2, s => 2] = -1.0
  Op[s' => 3, s => 3] = -1.0
  return Op[s' => 4, s => 4] = +1.0
end

function ITensors.op!(Op::ITensor, ::OpName"Fup", ::SiteType"Electron", s::Index)
  Op[s' => 1, s => 1] = +1.0
  Op[s' => 2, s => 2] = -1.0
  Op[s' => 3, s => 3] = +1.0
  return Op[s' => 4, s => 4] = -1.0
end

function ITensors.op!(Op::ITensor, ::OpName"Fdn", ::SiteType"Electron", s::Index)
  Op[s' => 1, s => 1] = +1.0
  Op[s' => 2, s => 2] = +1.0
  Op[s' => 3, s => 3] = -1.0
  return Op[s' => 4, s => 4] = -1.0
end

function ITensors.op!(Op::ITensor, ::OpName"Sz", ::SiteType"Electron", s::Index)
  Op[s' => 2, s => 2] = +0.5
  return Op[s' => 3, s => 3] = -0.5
end

function ITensors.op!(Op::ITensor, ::OpName"Sᶻ", st::SiteType"Electron", s::Index)
  return op!(Op, OpName("Sz"), st, s)
end

function ITensors.op!(Op::ITensor, ::OpName"Sx", ::SiteType"Electron", s::Index)
  Op[s' => 2, s => 3] = 0.5
  return Op[s' => 3, s => 2] = 0.5
end

function ITensors.op!(Op::ITensor, ::OpName"Sˣ", st::SiteType"Electron", s::Index)
  return op!(Op, OpName("Sx"), st, s)
end

function ITensors.op!(Op::ITensor, ::OpName"S+", ::SiteType"Electron", s::Index)
  return Op[s' => 2, s => 3] = 1.0
end

function ITensors.op!(Op::ITensor, ::OpName"S⁺", st::SiteType"Electron", s::Index)
  return op!(Op, OpName("S+"), st, s)
end
function ITensors.op!(Op::ITensor, ::OpName"Sp", st::SiteType"Electron", s::Index)
  return op!(Op, OpName("S+"), st, s)
end
function ITensors.op!(Op::ITensor, ::OpName"Splus", st::SiteType"Electron", s::Index)
  return op!(Op, OpName("S+"), st, s)
end

function ITensors.op!(Op::ITensor, ::OpName"S-", ::SiteType"Electron", s::Index)
  return Op[s' => 3, s => 2] = 1.0
end

function ITensors.op!(Op::ITensor, ::OpName"S⁻", st::SiteType"Electron", s::Index)
  return op!(Op, OpName("S-"), st, s)
end
function ITensors.op!(Op::ITensor, ::OpName"Sm", st::SiteType"Electron", s::Index)
  return op!(Op, OpName("S-"), st, s)
end
function ITensors.op!(Op::ITensor, ::OpName"Sminus", st::SiteType"Electron", s::Index)
  return op!(Op, OpName("S-"), st, s)
end

ITensors.has_fermion_string(::OpName"Cup", ::SiteType"Electron") = true
ITensors.has_fermion_string(::OpName"Cdagup", ::SiteType"Electron") = true
ITensors.has_fermion_string(::OpName"Cdn", ::SiteType"Electron") = true
ITensors.has_fermion_string(::OpName"Cdagdn", ::SiteType"Electron") = true


# ExoplanetsSysSim/src/summary_plots.jl
## (c) 2015 Darin Ragozzine

try
if !(Pkg.installed("PyPlot") >= v"2.1.0")  # version chosen arbitrarily
  error("You need to upgrade the Optim package for this example.")
end
catch
  error("You need to install the Optim package for this example.")
end

verbose = true

export plot_all_summary_stats

if length(ARGS)<2
   warn("usage:  syssim_plots.jl param_file comp_file")
   println("ASSUMING!!: param_file = demo_param.in AND comp_file = demo_ss.out")

   ## This hack allows this command to be run from inside the REPL or at the
   ## command line, though Julia understandably doesn't like it.
   ARGS=["demo_param.in","demo_ss.out"]
   # exit(255)
end

param_file = convert(String,ARGS[1])
if !isreadable(param_file) || !isfile(param_file)
   warn(string("Can't read ",param_file))
   exit(255)
end

comp_file = convert(String,ARGS[2])
if !isreadable(comp_file) || !isfile(comp_file)
   warn(string("Can't read ",comp_file))
   exit(255)
end

if verbose println("# Loading ExoplanetsSysSim module."); end
using ExoplanetsSysSim
if verbose println("# Loading PyPlot Module; that's how we'll do plots."); end
using PyPlot
pygui(false)  # Tell PyPlot not to use the GUI, but to return objects

if verbose println("# Reading parameter file ", param_file,"."); end
include(param_file)

if verbose println("# Calling setup_sim_param( ", ARGS[3:end], " )."); end
sim_param = setup_sim_param( convert(Array{String,1},ARGS[3:end]) )
if verbose println("# Active parameter values: ", make_vector_of_sim_param(sim_param) ); end

if verbose println("# Loading summary statistics from ", comp_file,"."); end
ss_comp = load_summary_stats(comp_file)

# Maybe?
type PlotSummaryStatistics
  plotnames::Dict{String,Any}  	# For storing names of plots
  plotfilenames::Dict{String,Any}  # For storing filenames of plots
end
function PlotSummaryStatistics()
  PlotSummaryStatistics( Dict{String,Any}(), Dict{String,Any}() )
end

function plot_all_summary_stats(ss::CatalogSummaryStatistics,
      sim_param::SimParam, verbose::Bool = false)
## TODO: Figure out the best filename to output to (user defined?) and implement.

# Print P vs. depth log-log scatter plot
if haskey(ss.stat,"P list") && haskey(ss.stat,"depth list")
	figure()
	## Some searching and outputting directly to a file without popping
	## up a figure window is not obvious, at least using PyPlot. Could use Gadfly.
	## TODO: Allow to output directly to a file without popping up a window
	loglog(ss.stat["P list"],ss.stat["depth list"],".")
    savefig("P_vs_depth.png")
else
	if verbose println("Cant plot P vs. depth because these are not found."); end
end #plot P vs. depth

# Plot multiplicity histogram in semilogy histogram-y scatter plot
if haskey(ss.stat,"num_sys_tranets")
	figure()
	semilogy(ss.stat["num_sys_tranets"],drawstyle="steps")
    savefig("num_sys_tranets.png")
else
	if verbose println("Cant plot multiplicity histogram because not found."); end
end # plot multiplicity histogram

end # plot_all_summary_stats

plot_all_summary_stats(ss_comp, sim_param, verbose)


"""
    $(SIGNATURES)

Perform the nonlinear numerical operations to approximate the convolution with a particular user defined likelihood function (conditional), which as been prepared in the `frl` object.  This function uses root finding to enforce a non-linear function constraint.

Notes:
- remember this is a deepcopy of original sfidx, since we are generating a proposal distribution and not directly replacing the existing variable belief estimate

Future work:
- once Threads.@threads have been optmized JuliaLang/julia#19967, also see area4 branch
- improve handling of n and particleidx, especially considering future multithreading support
"""
function approxConvOnElements!( ccwl::Union{CommonConvWrapper{F},
                                            CommonConvWrapper{Mixture{N_,F,S,T}}},
                                elements::Union{Vector{Int}, UnitRange{Int}},
                                ::Type{<:MultiThreaded},
                                _slack=nothing ) where {N_,F<:AbstractRelative,S,T}
  #
  Threads.@threads for n in elements
    # ccwl.thrid_ = Threads.threadid()
    ccwl.cpt[Threads.threadid()].particleidx = n

    # ccall(:jl_, Nothing, (Any,), "starting loop, thrid_=$(Threads.threadid()), partidx=$(ccwl.cpt[Threads.threadid()].particleidx)")
    _solveCCWNumeric!( ccwl, _slack=_slack)
  end
  nothing
end

function approxConvOnElements!( ccwl::Union{CommonConvWrapper{F},
                                            CommonConvWrapper{Mixture{N_,F,S,T}}},
                                elements::Union{Vector{Int}, UnitRange{Int}},
                                ::Type{<:SingleThreaded},
                                _slack=nothing ) where {N_,F<:AbstractRelative,S,T}
  #
  for n in elements
    ccwl.cpt[Threads.threadid()].particleidx = n
    _solveCCWNumeric!( ccwl, _slack=_slack)
  end
  nothing
end

function approxConvOnElements!( ccwl::Union{CommonConvWrapper{F},
                                            CommonConvWrapper{Mixture{N_,F,S,T}}},
                                elements::Union{Vector{Int}, UnitRange{Int}},
                                _slack=nothing )  where {N_,F<:AbstractRelative,S,T}
  #
  approxConvOnElements!(ccwl, elements, ccwl.threadmodel, _slack)
end

"""
    $SIGNATURES

Control the amount of entropy to add to null-hypothesis in multihypo case.

Notes:
- Basically calculating the covariance (with a bunch of assumptions TODO, fix)
- FIXME, Currently only supports Euclidean domains.
- FIXME, allow particle subpopulations instead of just all of a variable
"""
function calcVariableDistanceExpectedFractional(ccwl::CommonConvWrapper,
                                                sfidx::Integer,
                                                certainidx::AbstractVector{<:Integer};
                                                kappa::Real=3.0  )
  #
  if sfidx in certainidx
    # msst_ = sqrt(calcCovarianceBasic(getManifold(ccwl.vartypes[sfidx]), ccwl.params[sfidx]))
    msst_ = calcStdBasicSpread(ccwl.vartypes[sfidx](), ccwl.params[sfidx])
    return kappa*msst_
  end
  # @assert !(sfidx in certainidx) "null hypo distance does not work for sfidx in certainidx"

  # get mean of all fractional variables
  # ccwl.params::Vector{Vector{P}}
  uncertainidx = setdiff(1:length(ccwl.params), certainidx)
  dists = zeros(length(uncertainidx)+length(certainidx))

  dims = manifold_dimension(getManifold(ccwl.vartypes[sfidx]))

  uncMeans = zeros(dims, length(uncertainidx))
  for (count,i) in enumerate(uncertainidx)
    u = mean(getManifold(ccwl.vartypes[i]), ccwl.params[i])
    uncMeans[:,count] .= getCoordinates(ccwl.vartypes[i], u)
  end
  count = 0

  refMean = getCoordinates(ccwl.vartypes[sfidx], mean(getManifold(ccwl.vartypes[sfidx]), ccwl.params[sfidx]))

  # calc for uncertain and certain
  for i in uncertainidx
    count += 1
    dists[count] = norm(refMean - uncMeans[:,count])
  end
  # also check distance to certainidx for general scale reference (workaround heuristic)
  for cidx in certainidx
    count += 1
    cerMeanPnt = mean(getManifold(ccwl.vartypes[cidx]), ccwl.params[cidx])
    cerMean = getCoordinates(ccwl.vartypes[cidx], cerMeanPnt)
    dists[count] = norm(refMean[1:dims] - cerMean[1:dims])
  end

  push!(dists, 1e-2)
  return kappa*maximum(dists)
end

# Add entrypy on a point in `points` on manifold M, only on dimIdx if in p
function addEntropyOnManifold!( M::ManifoldsBase.AbstractManifold,
                                    points::Union{<:AbstractVector{<:Real},SubArray},
                                    dimIdx::AbstractVector,
                                    spreadDist::Real,
                                    p::Union{Colon, <:AbstractVector}=: )
  #
  if length(points) == 0
    return nothing
  end

  # preallocate
  T = number_eltype(points[1])
  Xc = zeros(T, manifold_dimension(M))
  X = get_vector(M, points[1], Xc, DefaultOrthogonalBasis())

  for idx in 1:length(points)
    # build tangent coordinate random
    for dim in dimIdx
      if (p === :) || dim in p
        Xc[dim] = spreadDist*(rand()-0.5)
      end
    end
    # update tangent vector X
    get_vector!(M, X, points[idx], Xc, DefaultOrthogonalBasis())
    #update point
    exp!(M, points[idx], points[idx], X)
    # points[idx] = exp(M, points[idx], X)

  end
  #
  nothing
end

"""
    $(SIGNATURES)

Common function to compute across a single user defined multi-hypothesis ambiguity per factor.
This function dispatches both `AbstractRelativeRoots` and `AbstractRelativeMinimize` factors.

DevNotes
- Future combo with `_calcIPCRelative`
"""
function computeAcrossHypothesis!(ccwl::Union{<:CommonConvWrapper{F},
                                              <:CommonConvWrapper{Mixture{N_,F,S,T}}},
                                  hyporecipe::NamedTuple,
                                  sfidx::Int,
                                  maxlen::Int,
                                  mani::ManifoldsBase.AbstractManifold; # maniAddOps::Tuple;
                                  spreadNH::Real=5.0,
                                  inflateCycles::Int=3,
                                  skipSolve::Bool=false,
                                  testshuffle::Bool=false,
                                  _slack=nothing ) where {N_,F<:AbstractRelative,S,T}
  #
  count = 0
  # transition to new hyporecipe approach
  allelements = hyporecipe.allelements
  activehypo  = hyporecipe.activehypo
  certainidx  = hyporecipe.certainidx

  cpt_ = ccwl.cpt[Threads.threadid()]

  @assert ccwl.certainhypo == hyporecipe.certainidx "expected hyporecipe.certainidx to be the same as cached in ccw"
  for (hypoidx, vars) in activehypo
    count += 1

    # now do hypothesis specific
    if sfidx in certainidx && hypoidx != 0 || hypoidx in certainidx || hypoidx == sfidx
      # hypo case hypoidx, sfidx = $hypoidx, $sfidx
      for i in 1:Threads.nthreads()  ccwl.cpt[i].activehypo = vars; end

      addEntr = view(ccwl.params[sfidx], allelements[count])
      # dynamic estimate with user requested speadNH of how much noise to inject (inflation or nullhypo)
      spreadDist = calcVariableDistanceExpectedFractional(ccwl, sfidx, certainidx, kappa=ccwl.inflation)

      # do proposal inflation step, see #1051
      # consider duplicate convolution approximations for inflation off-zero
      # ultimately set by dfg.params.inflateCycles
      for iflc in 1:inflateCycles
        addEntropyOnManifold!(mani, addEntr, 1:getDimension(mani), spreadDist, ccwl.partialDims)
        # no calculate new proposal belief on kernels `allelements[count]`
          _checkErrorCCWNumerics(ccwl, testshuffle)
        skipSolve ? @warn("skipping numerical solve operation") : approxConvOnElements!(ccwl, allelements[count], _slack)
      end
    elseif hypoidx != sfidx && hypoidx != 0
      # snap together case
      # multihypo, take other value case
      # sfidx=2, hypoidx=3:  2 should take a value from 3
      # sfidx=3, hypoidx=2:  3 should take a value from 2
      # DEBUG sfidx=2, hypoidx=1 -- bad when do something like multihypo=[0.5;0.5] -- issue 424
      # ccwl.params[sfidx][:,allelements[count]] = view(ccwl.params[hypoidx],:,allelements[count])
        # NOTE make alternative case only operate as null hypo
        addEntr = view(ccwl.params[sfidx], allelements[count])
        # dynamic estimate with user requested speadNH of how much noise to inject (inflation or nullhypo)
        spreadDist = calcVariableDistanceExpectedFractional(ccwl, sfidx, certainidx, kappa=spreadNH)
        addEntropyOnManifold!(mani, addEntr, 1:getDimension(mani), spreadDist)

    elseif hypoidx == 0
      # basically do nothing since the factor is not active for these allelements[count]
      # inject more entropy in nullhypo case
      # add noise (entropy) to spread out search in convolution proposals
      addEntr = view(ccwl.params[sfidx], allelements[count])
      # dynamic estimate with user requested speadNH of how much noise to inject (inflation or nullhypo)
      spreadDist = calcVariableDistanceExpectedFractional(ccwl, sfidx, certainidx, kappa=spreadNH)
      # # make spread (1σ) equal to mean distance of other fractionals
      addEntropyOnManifold!(mani, addEntr, 1:getDimension(mani), spreadDist)
    else
      error("computeAcrossHypothesis -- not dealing with multi-hypothesis case correctly")
    end
  end
  nothing
end
  # elseif hypoidx == sfidx
  #   # multihypo, do conv case, hypoidx == sfidx
  #   ah = sort(union([sfidx;], certainidx))
  #   @assert norm(ah - vars) < 1e-10
  #   for i in 1:Threads.nthreads()  ccwl.cpt[i].activehypo = ah; end
  #   approxConvOnElements!(ccwl, allelements[count])

# TODO what about nullhypo in recipe (when .mhidx[smpid]==0)?
# TODO figure out how best to combine with computeAcrossHypothesis!
function _calcIPCRelative(Xi::AbstractVector{<:DFGVariable},
                          ccwl::CommonConvWrapper,
                          hyporecipe::NamedTuple,
                          sfidx::Integer,
                          smpid::Integer=findfirst(x->x != 0, hyporecipe.mhidx)
                          )
  #
  @assert hyporecipe.activehypo[1][1] === 0 "expected 0-hypo case in hyporecipe.activehypo, to get variable hypo mask for relative partial propagation calculations."
  @assert hyporecipe.mhidx[smpid] !== 0 "_calcIPCRelative does not yet handle nullhypo gradients, try alternative hypo (smpid=$smpid), available hypos are hyporecipe.mhidx=$(hyporecipe.mhidx)"

  # select only the active variables in case of multihypo
  # @show smpid
  # @show hyporecipe.mhidx
  # @show hyporecipe.activehypo
  # NOTE +1 bc first element in .activehypo is nullhypo case, e.g. `(0,[1;])`
  _selhypo = hyporecipe.mhidx[smpid] + 1
  activehypo = hyporecipe.activehypo[_selhypo]
  activeids = activehypo[2]
  # solvefor index without the fractional variables
  active_mask = (x->x in activeids).(1:length(Xi))
  sfidx_active = sum(active_mask[1:sfidx])

  # build a view to the decision variable memory
  activeParams = view(ccwl.params, activeids)
  activeVars = Xi[active_mask]

  # assume gradients are just done for the first sample values
  # error("Possible issue, a factor has one manifold and attached variables have different manifolds.  Make sure the plumbing respects that.")
  @show typeof(ccwl.usrfnc!)
  @show sfidx
  # @show getLabel.(Xi)
  @show getLabel.(activeVars)
  @show getVariableType.(activeVars)
  # @show _getindextuple(ccwl.measurement, smpid)
  meas_pts = tuple((_getindextuple(ccwl.measurement, smpid))..., (getindex.(activeParams, smpid))...)
  # @show meas_pts
  #
  ipc = if ccwl._gradients === nothing
    ones(getDimension(activeVars[sfidx_active]))
  else
    ipc_ = Pair[]
    # get infoPerCoord from all variables
    for (vid,var) in enumerate(activeVars)
      # set all other variables infoPerCoord values
      getLabel(var) != getLabel(activeVars[sfidx_active]) ? nothing : continue
      push!(ipc_, vid=>ones(getDimension(var)))
    end
    # update the gradients at current point estimates
    # meas_pts =
    ccwl._gradients(meas_pts...)
    # do perturbation check
    # @show ipc_
    allipc = calcPerturbationFromVariable(ccwl._gradients, ipc_)
    allipc[sfidx_active]
  end

  @show ipc
  # FIXME REMOVE, overwrite with defauls during dev
  # fill!(ipc, 1.0)

  return ipc
end

"""
    $(SIGNATURES)

Multiple dispatch wrapper for `<:AbstractRelativeRoots` types, to prepare and execute the general approximate convolution with user defined factor residual functions.  This method also supports multihypothesis operations as one mechanism to introduce new modality into the proposal beliefs.

Planned changes will fold null hypothesis in as a standard feature and no longer appear as a separate `InferenceVariable`.
"""
function evalPotentialSpecific( Xi::AbstractVector{<:DFGVariable},
                                ccwl::CommonConvWrapper{T},
                                solvefor::Symbol,
                                T_::Type{<:AbstractRelative},
                                measurement::AbstractVector=Tuple[];
                                needFreshMeasurements::Bool=true,
                                solveKey::Symbol=:default,
                                N::Int= 0<length(measurement) ? length(measurement) : maximum(Npts.(getBelief.(Xi, solveKey))),
                                spreadNH::Real=3.0,
                                inflateCycles::Int=3,
                                nullSurplus::Real=0,
                                dbg::Bool=false,
                                skipSolve::Bool=false,
                                _slack=nothing  ) where {T <: AbstractFactor}
  #

  # Prep computation variables
  # NOTE #1025, should FMD be built here...
  sfidx, maxlen, mani = _updateCCW!(ccwl, Xi, solvefor, N; needFreshMeasurements, solveKey)
  # check for user desired measurement values
  if 0 < length(measurement)
    ccwl.measurement = measurement
  end

  # Check which variables have been initialized
  isinit = map(x->isInitialized(x), Xi)

  # assemble how hypotheses should be computed
  # nullSurplus see #1517
  runnullhypo = maximum((ccwl.nullhypo,nullSurplus))
  hyporecipe = _prepareHypoRecipe!(ccwl.hypotheses, maxlen, sfidx, length(Xi), isinit, runnullhypo )

  # get manifold add operations
  # TODO, make better use of dispatch, see JuliaRobotics/RoME.jl#244
  # addOps, d1, d2, d3 = buildHybridManifoldCallbacks(manis)
  mani = getManifold(getVariableType(Xi[sfidx]))

  # perform the numeric solutions on the indicated elements
  # FIXME consider repeat solve as workaround for inflation off-zero
  # NOTE alternate use of ccwl.certainidx to hyporecipe, certainidx = ccwl.certainhypo
  computeAcrossHypothesis!( ccwl, hyporecipe,
                            sfidx, maxlen, mani, spreadNH=spreadNH,
                            inflateCycles=inflateCycles, skipSolve=skipSolve,
                            _slack=_slack )
  #
  # FIXME do info per coord
  # ipc_ = _calcIPCRelative(Xi, ccwl, hyporecipe, sfidx)
  ipc = ones(getDimension(Xi[sfidx]))
  if isPartial(ccwl)
    # FIXME this is a workaround until better _calcIPCRelative can be used
    msk_ = setdiff(1:length(ipc), ccwl.usrfnc!.partial)
    for _i in msk_
      ipc[_i] = 0.0
    end
  end

  # return the found points, and info per coord
  return ccwl.params[ccwl.varidx], ipc
end

# TODO `measurement` might not be properly wired up yet
# TODO consider 1051 here to inflate proposals as general behaviour
function evalPotentialSpecific( Xi::AbstractVector{<:DFGVariable},
                                ccwl::CommonConvWrapper{T},
                                solvefor::Symbol,
                                T_::Type{<:AbstractPrior},
                                measurement::AbstractVector=Tuple[];
                                needFreshMeasurements::Bool=true,
                                solveKey::Symbol=:default,
                                N::Int=length(measurement),
                                dbg::Bool=false,
                                spreadNH::Real=3.0,
                                inflateCycles::Int=3,
                                nullSurplus::Real=0,
                                skipSolve::Bool=false,
                                _slack=nothing ) where {T <: AbstractFactor}
  #
  # setup the partial or complete decision variable dimensions for this ccwl object
  # NOTE perhaps deconv has changed the decision variable list, so placed here during consolidation phase
  _setCCWDecisionDimsConv!(ccwl)

  # FIXME, NEEDS TO BE CLEANED UP AND WORK ON MANIFOLDS PROPER
  fnc = ccwl.usrfnc!
  sfidx = findfirst(getLabel.(Xi) .== solvefor)
  # sfidx = 1 #  WHY HARDCODED TO 1??
  solveForPts = getVal(Xi[sfidx], solveKey=solveKey)
  nn = maximum([N; calcZDim(ccwl); length(solveForPts); length(ccwl.params[sfidx])])  # length(measurement[1])

  # FIXME better standardize in-place operations (considering solveKey)
  if needFreshMeasurements
    cf = CalcFactor( ccwl )
    newMeas = sampleFactor(cf, nn)
    ccwl.measurement = newMeas
  end

  # Check which variables have been initialized
  # TODO not sure why forcing to Bool vs BitVector
  isinit::Vector{Bool} = Xi .|> isInitialized .|> Bool
  # nullSurplus see #1517
  runnullhypo = maximum((ccwl.nullhypo,nullSurplus))
  hyporecipe = _prepareHypoRecipe!(ccwl.hypotheses, nn, sfidx, length(Xi), isinit, runnullhypo )

  # get solvefor manifolds, FIXME ON FIRE, upgrade to new Manifolds.jl
  mani = getManifold(Xi[sfidx])
  # two cases on how to use the measurement
  nhmask = hyporecipe.mhidx .== 0
  ahmask = hyporecipe.mhidx .== 1
  # generate nullhypo samples
  # inject lots of entropy in nullhypo case
  # make spread (1σ) equal to mean distance of other fractionals
  # FIXME better standardize in-place operations (considering solveKey)
  addEntr = if length(solveForPts) == nn
    deepcopy(solveForPts)
  else
    ret = typeof(solveForPts)(undef, nn)
    for i in 1:length(solveForPts)
      ret[i] = solveForPts[i]
    end
    for i in (length(solveForPts)+1):nn
      ret[i] = getPointIdentity(getVariableType(Xi[sfidx]))
    end
    ret
  end

  # view on elements marked for nullhypo
  addEntrNH = view(addEntr, nhmask)
  # spreadDist = spreadNH*sqrt(calcCovarianceBasic(mani, addEntr))
  spreadDist = spreadNH*calcStdBasicSpread(getVariableType(Xi[sfidx]), addEntr)
  # partials are treated differently
  ipc = if !isPartial(ccwl) #ccwl.partial
    # TODO for now require measurements to be coordinates too
    # @show typeof(ccwl.measurement[1])
    for m in (1:length(addEntr))[ahmask]
      # FIXME, selection for all measurement::Tuple elements
      # @info "check broadcast" ccwl.usrfnc! addEntr[m] ccwl.measurement[1][m]
      setPointsMani!(addEntr[m], ccwl.measurement[m])
      # addEntr[m] = ccwl.measurement[m][1]
    end
    # ongoing part of RoME.jl #244
    addEntropyOnManifold!(mani, addEntrNH, 1:getDimension(mani), spreadDist)
    # do info per coords
    ones(getDimension(Xi[sfidx]))
  else
    # FIXME but how to add partial factor info only on affected dimensions fro general manifold points?
    pvec = [fnc.partial...]

    if !hasmethod(getManifold, (typeof(fnc),))
      @debug "No method getManifold for $(typeof(fnc)), using getManifoldPartial"
    end

    # active hypo that receives the regular measurement information
    for m in (1:length(addEntr))[ahmask]
      # addEntr is no longer in coordinates, these are now general manifold points!!
      # for (i,dimnum) in enumerate(fnc.partial)
        # FIXME, need ability to replace partial points
        partialCoords = ccwl.partialDims

        #FIXME check if getManifold is defined otherwise fall back to getManifoldPartial, JT: I would like to standardize to getManifold
        if hasmethod(getManifold, (typeof(fnc),))
          Msrc = getManifold(fnc)
          # TODO workaround until partial manifold approach is standardized, see #1492
          asPartial = isPartial(fnc) && manifold_dimension(Msrc) < manifold_dimension(mani)
          setPointPartial!(mani, addEntr[m], Msrc, ccwl.measurement[m], partialCoords, asPartial)
        else
          Msrc, = getManifoldPartial(mani, partialCoords)
          setPointPartial!(mani, addEntr[m], Msrc, ccwl.measurement[m], partialCoords)
        end
        # addEntr[m][dimnum] = ccwl.measurement[1][m][i]
      # end
    end
    # null hypo mask that needs to be perturbed by "noise"
    addEntrNHp = view(addEntr, nhmask)
    # ongoing part of RoME.jl #244
    addEntropyOnManifold!(mani, addEntrNHp, 1:getDimension(mani), spreadDist, pvec)
    # do info per coords
    ipc_ = zeros(getDimension(Xi[sfidx]))
    ipc_[pvec] .= 1.0
    ipc_
  end

  # check partial is easy as this is a prior
  return addEntr, ipc
end

function evalPotentialSpecific( Xi::AbstractVector{<:DFGVariable},
                                ccwl::CommonConvWrapper{Mixture{N_,F,S,T}},
                                solvefor::Symbol,
                                measurement::AbstractVector=Tuple[];
                                kw... ) where {N_,F<:AbstractFactor,S,T}
  #
  evalPotentialSpecific(Xi,
                        ccwl,
                        solvefor,
                        F,
                        measurement;
                        kw... )
end

function evalPotentialSpecific( Xi::AbstractVector{<:DFGVariable},
                                ccwl::CommonConvWrapper{F},
                                solvefor::Symbol,
                                measurement::AbstractVector=Tuple[];
                                kw... ) where {F <: AbstractFactor}
  #
  evalPotentialSpecific(Xi,
                        ccwl,
                        solvefor,
                        F,
                        measurement;
                        kw... )
end

"""
    $(SIGNATURES)

Single entry point for evaluating factors from factor graph, using multiple dispatch to locate the correct `evalPotentialSpecific` function.
"""
function evalFactor(dfg::AbstractDFG,
                    fct::DFGFactor,
                    solvefor::Symbol,
                    measurement::AbstractVector=Tuple[];
                    needFreshMeasurements::Bool=true,
                    solveKey::Symbol=:default,
                    N::Int=length(measurement),
                    inflateCycles::Int=getSolverParams(dfg).inflateCycles,
                    nullSurplus::Real=0,
                    dbg::Bool=false,
                    skipSolve::Bool=false,
                    _slack=nothing  )
  #

  ccw = _getCCW(fct)
  # TODO -- this build up of Xi is excessive and could happen at addFactor time
  variablelist = getVariableOrder(fct)
  Xi = getVariable.(dfg, variablelist)

  # setup operational values before compute (likely to be refactored)
  for i in 1:Threads.nthreads()
    ccw.cpt[i].factormetadata.variablelist = variablelist
    ccw.cpt[i].factormetadata.solvefor = solvefor
  end

  return evalPotentialSpecific( Xi, ccw, solvefor, measurement; needFreshMeasurements,
                                solveKey, N, dbg, spreadNH=getSolverParams(dfg).spreadNH,
                                inflateCycles, nullSurplus, skipSolve, _slack  )
  #
end

"""
    $SIGNATURES

Perform factor evaluation to resolve the "solve for" variable of a factor.
This temporary function can be run without passing a factor graph object, but will internally allocate a new temporary new one.
Alternatively, the factor graph used for calculations can be passed in via the keyword `tfg`, hence the function name bang.

Notes
- `TypeParams_args::Vector{Tuple{InferenceVariable, P}}
- idea is please find best e.g. `b`, given `f(z,a,b,c)` either by roots or minimize (depends on definition of `f`)
- `sfidx::Int` is the solve for index, assuming `getVariableOrder(fct)`.

Example
```julia
B = _evalFactorTemporary!(EuclidDistance, (ContinuousScalar, ContinuousScalar), 2, ([10;],), ([0.],[9.5]) )
# should return `B = 10`
```

Related

[`evalFactor`](@ref), [`calcFactorResidual`](@ref), [`testFactorResidualBinary`](@ref), [`solveFactorParameteric`](@ref), [`approxConvBelief`](@ref)
"""
function _evalFactorTemporary!( fct::AbstractFactor,
                                varTypes::Tuple,
                                sfidx::Int,  # solve for index, assuming variable order for fct
                                measurement::AbstractVector,
                                pts::Tuple;
                                tfg::AbstractDFG=initfg(),
                                solveKey::Symbol=:default,
                                newFactor::Bool=true,
                                _slack=nothing,
                                buildgraphkw... )
  #

  # build up a temporary graph in dfg
  _, _dfgfct = IIF._buildGraphByFactorAndTypes!(fct, varTypes, pts; dfg=tfg, solveKey=solveKey, newFactor=newFactor, buildgraphkw...)

  # get label convention for which variable to solve for
  solvefor = getVariableOrder(_dfgfct)[sfidx]

  # do the factor evaluation
  sfPts, _ = evalFactor(tfg, _dfgfct, solvefor, measurement, needFreshMeasurements=false, solveKey=solveKey, inflateCycles=1, _slack=_slack )

  return sfPts
end

#


include(joinpath(pwd(), "models/visualize.jl"))
vis = Visualizer()
open(vis)
default_background!(vis)

# box
dir = "/home/taylor/Research/ContactControl.jl/src/dynamics/quadruped/payload_mesh/Box.obj"
obj = MeshFileObject(dir)
setobject!(vis[:box], obj)
tr = Translation(0.495, 0.0, 0.3)
rt = LinearMap(0.2 * RotZ(-0.5 * π) * RotX(1.0 * π))
tform = compose(tr, rt)
settransform!(vis[:box], tform)

include_model("simple_manipulator")

# arm 1
l1 = Cylinder(Point3f0(0, 0, 0),
	Point3f0(0, 0, 1.0),
	convert(Float32, 0.025))
setobject!(vis["l1_1"], l1, MeshPhongMaterial(color = RGBA(0, 0, 0, 1.0)))
l2 = Cylinder(Point3f0(0, 0, 0), Point3f0(0, 0, 1.0),
	convert(Float32, 0.025))
setobject!(vis["l2_1"], l2, MeshPhongMaterial(color = RGBA(0, 0, 0, 1.0)))

setobject!(vis["elbow_1"], Sphere(Point3f0(0),
	convert(Float32, 0.05)),
	MeshPhongMaterial(color = RGBA(0, 0, 0, 1.0)))
setobject!(vis["ee_1"], Sphere(Point3f0(0),
	convert(Float32, 0.05)),
	MeshPhongMaterial(color = RGBA(1.0, 165.0 / 255.0, 0.0, 1.0)))

p1 = [-0.99; -0.01; 0.0]
q1 = [0.41 * π; -0.635 * π]
p_mid = [kinematics_mid(model, q1)[1], 0.0, kinematics_mid(model, q1)[2]] + p1
p_ee = [kinematics_ee(model, q1)[1], 0.0, kinematics_ee(model, q1)[2]] + p1

settransform!(vis["l1_1"], cable_transform(p1, p_mid))
settransform!(vis["l2_1"], cable_transform(p_mid, p_ee))

settransform!(vis["elbow_1"], Translation(p_mid))
settransform!(vis["ee_1"], Translation(p_ee))

# arm 2
setobject!(vis["l1_2"], l1, MeshPhongMaterial(color = RGBA(0, 0, 0, 1.0)))
setobject!(vis["l2_2"], l2, MeshPhongMaterial(color = RGBA(0, 0, 0, 1.0)))

setobject!(vis["elbow_2"], Sphere(Point3f0(0),
	convert(Float32, 0.05)),
	MeshPhongMaterial(color = RGBA(0, 0, 0, 1.0)))
setobject!(vis["ee_2"], Sphere(Point3f0(0),
	convert(Float32, 0.05)),
	MeshPhongMaterial(color = RGBA(1.0, 165.0 / 255.0, 0.0, 1.0)))

p2 = [0.99; -0.01; 0.0]
q2 = [0.4 *  π; 0.825 * π]
p_mid = [kinematics_mid(model, q2)[1], 0.0, kinematics_mid(model, q2)[2]] + p2
p_ee = [kinematics_ee(model, q2)[1], 0.0, kinematics_ee(model, q2)[2]] + p2

settransform!(vis["l1_2"], cable_transform(p2, p_mid))
settransform!(vis["l2_2"], cable_transform(p_mid, p_ee))

settransform!(vis["elbow_2"], Translation(p_mid))
settransform!(vis["ee_2"], Translation(p_ee))

settransform!(vis["/Cameras/default"],
	   compose(Translation(0.0, -95.0, -1.0), LinearMap(RotY(0.0 * π) * RotZ(-π / 2.0))))
setprop!(vis["/Cameras/default/rotated/<object>"], "zoom", 60)

setobject!(vis["box1"], GeometryBasics.HyperRectangle(Vec(0.0, 0.0, 0.0),
		Vec(0.25, 0.5, 0.25)), MeshPhongMaterial(color = RGBA(1.0, 0.0, 0.0, 1.0)))
	settransform!(vis["box1"], Translation(0.51, -0.25, 0))

setobject!(vis["box2"], GeometryBasics.HyperRectangle(Vec(0.0, 0.0, 0.0),
		Vec(0.25, 0.5, 0.25)), MeshPhongMaterial(color = RGBA(1.0, 0.0, 0.0, 1.0)))
	settransform!(vis["box2"], Translation(-0.16, -0.25, 0))


module SynthControl

using Dates, DataFrames, Optim, LinearAlgebra, RecipesBase, Parameters
import CSV

export TreatmentPanel, SynthControlModel, fit!, isfitted, load_brexit, load_germany, load_basque, load_smoking

abstract type SCM end

"""
    TreatmentPanel

Holds the data to be used in fitting synthetic control models. The object is constructed
from a `DataFrame` which contains all outcome and covariate data. The constructor requires
passing information on which column in the `DataFrame` holds the subject and time period
identifiers, as well as information on the timing of treatment. Treatments are specified
as a `Pair` of either a `String` or `Symbol` identifying the unit treated, and a `Date` or
`Int` value, identifying the treatment period. Where treatments have an end point, the
treatment period os specified as a `Tuple` of either `Date` or `Int`, indicating the first
and last period of treatment. Where individual units have multiple treatment periods,
these are specified as a `Pair{Union{String, Symbol}, Vector{Union{Int, Date}}}` of treatment
unit identifier and vector of timings. Finally, single continuous, single periodic, and
multiple period treatments can be generalized to multiple treated units.

The following table provides an overview of the types of treatment pattern supported:

|                 |  Only starting point        |   Start and end point                     |   Multiple start & end points                       |
|-----------      |------------------------     |-------------------------                  |------------------------------                       |
| **one unit**         |  Pair{String, Date}         |   Pair{String, Tuple{Date, Date}}         |  Pair{String}, Vector{Tuple{Date, Date}}}           |
| **multiple units**  |  Vector{Pair{String, Date}} |   Vector{Pair{String, Tuple{Date, Date}}} |  Vector{Pair{String}, Vector{Tuple{Date, Date}}}}   |

Currently, only single treatment unit and continuous treatment is supported.

"""

@with_kw struct TreatmentPanel
    data::DataFrame
    outcome::Union{Symbol, String}
    id_var::Union{Symbol, String}
    t_var::Union{Symbol, String}
    treatment::Union{Pair{Union{Symbol, String}, Union{Date, Int64}}}
    y_pre::Vector{Float64}
    y_post::Vector{Float64}
    x_pre::Matrix{Float64}
    x_post::Matrix{Float64}
    no_comps::Int64
    no_treatment_periods::Int64
    no_pretreatment_periods::Int64
    comp_labels::Vector{Union{Symbol, String}}
end

function check_tp(outcome, id_var, t_var, treatment, data)

    treatᵢ, treatₜ = treatment

    !isnothing(outcome) || error(ArgumentError(
        "Please specify outcome, the nameof the column in your dataset holding the "*
        "outcome variable of interest in your dataset."
        ))
    !isnothing(id_var) || error(ArgumentError(
            "Please specify id_var, the name of the column in your data set holding the "*
            "identifier of your units of observation (panel dimension)."
            ))
    !isnothing(t_var) || error(ArgumentError(
            "Please specify t_var, the name of the column in your dataset holding the "*
            "time dimension."
        ))
    !isnothing(treatment) || error(ArgumentError(
        "Please specify treatment, the identifier of the treated unit(s) in your "*
        "dataset and associated start(s) of treatment."
        ))
    in(string(outcome), names(data)) || error(DomainError(
        "The specified outcome column, $outcome, was not found in your data set. Please "*
        "double check the name of the outcome column."
        ))
    in(treatᵢ, data[!, id_var]) || error(DomainError(
        "The specified treatment unit, $(treatᵢ), was not found in the specified id "*
        "column, $id_var. Please ensure that treatment unit and id column are correctly "*
        "specified."
        ))
    in(treatₜ, data[!, t_var]) || error(DomainError(
        "The specified treatment period, $(treatₜ), was not found in the specified time "*
        "column, $t_var. Please double check treatment time and time column."
        ))
    treatₜ != maximum(data[!, t_var]) || error(DomainError(
        "The specified treatment period, $(treatₜ), is equal to the last observation period "*
        "of your data set - no estimation of treatment effects is possible. Please double "*
        "check the treatment period."
        ))
    treatₜ < maximum(data[!, t_var]) || error(DomainError(
        "The specified treatment period, $(treatₜ), occurs after the last observation in your "*
        "data set, $(maximum(data[!, t_var])) - no estimation of treatment effects is possible. "*
        "Please double check specification of treatment period."
        ))
    # To add - checks around balance of panel?
end

function TreatmentPanel(treatment::T, data;
                        outcome = nothing, id_var = nothing, t_var = nothing) where T <: Pair

    check_tp(outcome, id_var, t_var, treatment, data)

    treatᵢ, treatₜ = treatment

    y_pre = data[(data[!, id_var] .== treatᵢ) .&
               (data[!, t_var] .< treatₜ), outcome]
    xs_pre = data[(data[!, id_var] .!= treatᵢ) .&
               (data[!, t_var] .< treatₜ), outcome]

    y_post = data[(data[!, id_var] .== treatᵢ) .&
               (data[!, t_var] .>= treatₜ), outcome]
    xs_post = data[(data[!, id_var] .!= treatᵢ) .&
               (data[!, t_var] .>= treatₜ), outcome]

    no_treatment_periods = length(unique(data[data[!, t_var] .>= treatₜ, t_var]))
    no_pretreatment_periods = length(y_pre)
    no_comps = length(unique(data[data[!, id_var] .!= treatᵢ, id_var]))

    comp_labels = unique(data[data[!, id_var] .!= treatᵢ, id_var])

    TreatmentPanel(
        data,
        outcome,
        id_var,
        t_var,
        treatment,
        y_pre,
        y_post,
        reshape(xs_pre, no_comps, no_pretreatment_periods),
        reshape(xs_post, no_comps, no_treatment_periods),
        no_comps,
        no_treatment_periods,
        no_pretreatment_periods,
        comp_labels
    )
end

"""
    SynthControlModel

Takes in a `TreatmentPanel` and holds the results of calling the `fit!` method, optimal
weights for comparator units and the resulting predicted outcome of the synthetic control
unit as well as the estimated treatment effect from comparing synthetic control and observed
outcome of the treated unit.

"""

mutable struct SynthControlModel <: SCM
    treatment_panel::TreatmentPanel
    w::Vector{Float64} # Vector of weights
    ŷ::Vector{Float64} # Vector of post-treatment predicted outcomes for synthetic unit
    δ̂::Vector{Float64} # Vector of estimated treatmet impacts
    p_test_res::Array{Float64, 2} # Matrix of outcomes for untreated under placebo
end

function SynthControlModel(tp::TreatmentPanel)
    @unpack y_pre, y_post, no_comps, no_treatment_periods = tp

    ŷ = zeros(length(y_pre) + length(y_post))
    w = zeros(no_comps)
    δ̂ = zeros(no_treatment_periods)
    p_test_res = zeros(length(ŷ), no_comps)

    SynthControlModel(tp, w, ŷ, δ̂, p_test_res)
end

"""
isfitted(s::SynthControlModel)

Check whether the `SynthControlModel` object `s` has been fitted.

# Examples
```julia-repl
julia> df = load_brexit();

julia> s = SynthControlModel(df, :country, :dateid, :realgdp, 86, "United Kingdom");

julia> isfitted(s)
false

julia> fit!(s);

julia> isfitted(s)
true

```
"""
isfitted(s::SynthControlModel) = sum(s.w) ≈ 0.0 ? false : true

"""
    fit!(s::SynthControlModel; placebo_test = false)

Fit the `SynthControlModel` `s` by finding the weights that minimize the distance between
the pre-treatment outcomes for the observational unit of interest and the weighted average
pre-treatment outcomes for unweighted units.

If `placebo_test = true` is supplied, additional placebo tests will be performed by using
every non-treated unit in the data set as the treated unit in turn and estimating the
treatment impact on this unit. Results are stored in the `p_test_res` field and can be
used as the basis for inference.
"""
function fit!(s::SynthControlModel; placebo_test = false)

    @unpack y_pre, y_post, x_pre, x_post, no_comps, no_pretreatment_periods = s.treatment_panel

    # Objective function
    # Hat tip to Mathieu Tanneau who suggested switching to a quadratic formulation
    obj(w) = dot(y_pre .- vec(w'*x_pre), y_pre .- vec(w'*x_pre)) + 1e6*(1.0 - sum(w))^2

    # Initial condition and bounds
    initial = [1/no_comps for _ in 1:no_comps]
    lower = zeros(no_comps)
    upper = ones(no_comps)

    # Get weights through optimization
    res = optimize(obj, lower, upper, initial)
    s.w = res.minimizer

    # Get estimates
    s.ŷ = vec(s.w'*[x_pre x_post])
    s.δ̂ = ([y_pre; y_post] .- s.ŷ)[no_pretreatment_periods+1:end]

    if placebo_test
        placebos = unique(s.treatment_panel.data[(s.treatment_panel.data[!,s.pid] .!= s.treat_id), s.pid])
        p = deepcopy(s)

        for n ∈ 1:p.no_comps
            # change treatment ID
            p.treat_id = placebos[n]
            # change outcome data
            p.y = p.data[p.data[!, p.pid] .== p.treat_id, p.outcome]
            p.comps = collect(reshape(p.data[(p.data[!,p.pid] .!= p.treat_id) .& (p.data[!,p.tid].<p.treat_t), p.outcome],
                    p.no_comps, p.no_pretreat_t)')
            fit!(p)
            s.p_test_res[:, n] = p.y .- p.ŷ
        end
    end
    return s
end

# Define plotting recipe
@recipe function f(s::SynthControlModel; kind = "overall")
    legend --> :topleft

    if kind == "overall"

        @series begin
            label --> s.treatment_panel.treatment[1]
            sort!(unique(s.treatment_panel.data[!, s.treatment_panel.t_var])), [s.treatment_panel.y_pre; s.treatment_panel.y_post]
        end

        @series begin
            label --> "Control"
            sort!(unique(s.treatment_panel.data[!, s.treatment_panel.t_var])), s.ŷ
        end

        @series begin
            label --> "Impact"
            seriestype := :bar
            seriesalpha := 0.5
            linecolor := "white"
            seriescolor := "darkgreen"
            sort!(unique(s.treatment_panel.data[!, s.treatment_panel.t_var]))[findfirst(x -> x >= s.treatment_panel.treatment[2],
                                    sort!(unique(s.treatment_panel.data[!, s.treatment_panel.t_var]))):end], s.δ̂
        end

    elseif kind == "diffplot"
        @series begin
            label --> s.treatment_panel.treatment[1]
            sort!(unique(s.treatment_panel.data[!, s.treatment_panel.t_var])), [s.treatment_panel.y_pre; s.treatment_panel.y_post] .- s.ŷ
        end

        @series begin
            label --> ""
            seriestype := :scatter
            seriescolor := 1
            markerstrokecolor := "white"
            sort!(unique(s.treatment_panel.data[!, s.treatment_panel.t_var])), [s.treatment_panel.y_pre; s.treatment_panel.y_post] .- s.ŷ
        end
    end
        @series begin
            label --> "Intervention"
            seriestype := :vline
            linestyle := :dash
            seriescolor := "black"
            xguide := "Time"
            yguide := "Outcome"
            [s.treatment_panel.treatment[2]]
        end

     nothing
end

# Data loading for example
"""
    load_brexit()

Returns a `DataFrame` with quarterly GDP data on 30 OECD countries, borrowed from the
analysis of the effect of Brexit on UK GDP undertaken by the [Centre for European Reform](
https://www.cer.eu/insights/cost-brexit-june-2019)
"""
function load_brexit()
    CSV.read(joinpath(dirname(@__FILE__),"..","data","brexit.csv"), DataFrame)
end

function load_germany()
    CSV.read(joinpath(dirname(@__FILE__),"..","data","germany.csv"), DataFrame)
end

function load_basque()
    CSV.read(joinpath(dirname(@__FILE__),"..","data","basque.csv"), DataFrame)
end

function load_smoking()
    CSV.read(joinpath(dirname(@__FILE__),"..","data","smoking.csv"), DataFrame)
end

# Pretty printing
import Base.show

function show(io::IO, ::MIME"text/plain", s::SynthControlModel)

    println(io, "\nSynthetic Control Model")
    println(io, "\tOutcome variable: ", s.treatment_panel.outcome)
    println(io, "\tTime dimension: ",string(s.treatment_panel.t_var)," with ",
            string(length(unique(s.treatment_panel.data[!,s.treatment_panel.t_var]))), " unique values")
    println(io, "\tTreatment period: ", string(s.treatment_panel.treatment[2]))
    println(io, "\tID variable: ",string(s.treatment_panel.id_var), " with ",
            string(length(unique(s.treatment_panel.data[!,s.treatment_panel.id_var]))), " unique values")
    println(io, "\tTreatment ID: ",string(s.treatment_panel.treatment[1]))

    if isfitted(s)
      println(io, "\tModel is fitted")
      println(io, "\tImpact estimates: ",round.(s.δ̂, digits=3))
    else
      println(io, "\n\tModel is not fitted")
    end
end

end # module


# See https://github.com/ecmwf/cfgrib.jl for more information about CfGRIB
# You can install the HEAD version directly from github by hitting ] to enter package mode and then
#.   add https://github.com/ecmwf/cfgrib.jl/
# followed by backspace
using CfGRIB
using AxisArrays
using Printf
using Dates
using HTTP
using GZip

"""
readMRMS(fname)

Extracts MRMS data in grib format from a specified file and then extracts specific
data elements into CSV formatted files in a directory call csv (one file per location
with non-zero precipitation)

Saves data with detectable precipiation to folder ./csv,
based on truncated lat-lons.
Example: ./csv/48_70.csv
"""
function readMRMS(fname)
    fh = GZip.open(fname, "r")
    data = read(fh)
    io = open("data.grib2", "w")
    write(io, data)
    close(io)

    dataset = CfGRIB.DataSet("data.grib2")

    # Get datetime from dataset, but we could also parse the fname instead
    datetime = replace(string(unix2datetime(dataset.variables["time"].data)), "T" => " ")

    lats = dataset.variables["latitude"].data
    lons = dataset.variables["longitude"].data
    data = dataset.variables["paramId0"].data[:, :]

    files = Dict()

    # MRMS is 3500 x 7000
    for i = 1:3500, j = 1:7000
        lon = lons[j] - 360

        # Conterminous US data only
        if 25 <= lats[i] <= 50 && -125 <= lon <= -67
            # Exclude non-detectable precipitation to save space
            if data[j, i] > 0
                fn = "$(trunc(Int, lats[i]))_$(trunc(Int, -lon))"
                if !haskey(files, fn)
                    files[fn] = open("csv/$fn.csv", "w");
                    write(files[fn], "precipitation,date,lat,lon\n")
                end
                write(files[fn], "$(data[j, i]),$(datetime),$(@sprintf("%.4f", lats[i])),$(@sprintf("%.4f", lon))\r\n")
            end
        end
    end

    for file in values(files)
        close(file)
    end

    # not yet implemented
    # writeToPostgres()
end

"""
processMRMS()

Downloads data from the Iowa State archive of MRMS data.
Calls readMRMS() to extract and process the data.
Reruns after a 15 minute delay.
"""
function processMRMS()
    for date in Dates.Date(2021, 1, 1):Day(1):today()
        year = Dates.year(date)
        m2 = lpad(Dates.month(date), 2, "0")
        d2 = lpad(Dates.day(date),   2, "0")
        for hour in 0:23
            h2 = lpad(hour, 2, "0")
            fname = "MultiSensor_QPE_01H_Pass2_00.00_$year$m2$d2-$(h2)0000.grib2.gz"
            if !isfile("mrms/$fname")
                url = "https://mtarchive.geol.iastate.edu/$year/$m2/$d2/mrms/ncep/MultiSensor_QPE_01H_Pass2/$fname"
                try
                    download(url, "mrms/$fname")
                catch
                    rm("mrms/$fname") # file gets saved even if 404 error, so it must be removed
                    println("Couldn't download $url")
                    continue
                end

                println("Downloaded $fname")
                readMRMS("mrms/$fname")
            end
        end
    end
end

while true
    processMRMS()
    println("Pausing for 15 minutes")
    sleep(60 * 15)
end


#############################################
# Trie string approximate implementation
#############################################

# here we create a tree where each node is a letter, and the child nodes are the next letter in the string
# to find the UID for a string, we traverse the tree one letter at a time until we find the UID (or not!)

immutable TrieChar
  byte::UInt8
end
TrieChar() = TrieChar(0x00)

# Base.string(c::TrieChar) = ascii([c.byte])
Base.print(io::IO, c::TrieChar) = print(io, string(c))
Base.show(io::IO, c::TrieChar) = show(io, string(c))
Base.hash(c::TrieChar) = c.byte
==(c1::TrieChar, c2::TrieChar) = c1.byte == c2.byte

type TrieNode{T}
    children::Dict{TrieChar,TrieNode{T}}
    # parentChain::String
    # istail::Bool
    data::Nullable{T}  # note: this will be undefined except for tail nodes

    # TrieNode{T}(::Type{T}) = new(Dict{TrieChar,TrieNode{T}}(), "", false)
    # TrieNode{T}(::Type{T}, parentChain::String) = new(Dict{TrieChar,TrieNode{T}}(), parentChain, false)
    TrieNode() = new(Dict{TrieChar,TrieNode{T}}(), Nullable{T}())
end
# TrieNode{T}(::Type{T}) = TrieNode{T}(T)

function TrieNode{K}(::Type{K}, strings::Vector{String})
    rootNode = TrieNode{K}()
    for (i,s) in enumerate(strings)
        get!(rootNode, s, K(i))
    end
    rootNode
end

# function maketail!{T}(node::TrieNode{T}, data::T)
#     node.istail = true
#     node.data = data
#     return
# end

function Base.get{T<:Union{String,FixedLengthSymbol}, K}(rootNode::TrieNode{K}, s::T, default::K)
    local node::TrieNode{K}
    node = rootNode

    for i in 1:sizeof(T)
        b = getfield(s, i)
        if b == 0x00 || b == 0x20
            break
        end

        c = TrieChar(b)
        if haskey(node.children, c)
            node = node.children[TrieChar(b)]
        else
            return default
        end
        # node = get(node.children, TrieChar(b)) do
        #   return default
        # end
    end

    # if node.istail
    if isnull(node.data)
        return default
    else
        return get(node.data)
    end
end

function Base.get!{T<:Union{String,FixedLengthSymbol}, K}(rootNode::TrieNode{K}, s::T, default::K)
    node = rootNode
    for b in s.data
        # we're done if null char or space
        if b == 0x00 || b == 0x20
            break
        end

        # set node to be the next node in the tree (add one if needed)
        c = TrieChar(b)
        if haskey(node.children, c)
            node = node.children[c]
        else
            newnode = TrieNode{K}()
            node.children[c] = newnode
            node = newnode
        end

        # node = get!(node.children, TrieChar(b), TrieNode(K))

    end

    # add the tail
    if isnull(node.data)
        node.data = Nullable(default)
    end
    get(node.data)

    # maketail!(node, default)

    # node
end