GenericFunctions.jl

# TODO under development - experimenting with type to work with manifolds

## ======================================================================================
## Possible type piracy, but also standardizing to common API across repos
## ======================================================================================

DFG.getDimension(::Distributions.Uniform) = 1
DFG.getDimension(::Normal) = 1
DFG.getDimension(Z::MvNormal) = Z |> cov |> diag |> length

function DFG.getDimension(Z::FluxModelsDistribution)
  return if length(Z.outputDim) == 1
    Z.outputDim[1]
  else
    error(
      "can only do single index tensor at this time, please open an issue with Caesar.jl",
    )
  end
end
DFG.getDimension(Z::ManifoldKernelDensity) = getManifold(Z) |> getDimension
# TODO deprecate
DFG.getDimension(Z::BallTreeDensity) = Ndim(Z)

## ======================================================================================
## Generic manifold cost functions
## ======================================================================================
# """
#     $SIGNATURES
# Generic function that can be used in binary factors to calculate distance between points on Lie Groups with measurements.
# """
# function distancePoint2Point(M::SemidirectProductGroup, m, p, q)
#   q̂ = LieGroups.compose(M, p, m)
#   # return log(M, q, q̂)
#   return vee(M, q, log(M, q, q̂))
#   # return distance(M, q, q̂)
# end

# ::MeasurementOnTangent
function measurement_residual(G::AbstractLieGroup, X, p, q)
  X̂ = log(G, p, q)
  return vee(LieAlgebra(G), X - X̂) # TODO check sign with gradients, does not matter for cost so can't double check.
end

function prior_residual(G::AbstractLieGroup, m, p)
  #TODO should it be TₘM or TₚM?
  # Is the covariance that of the point m? If so, I would think it should be TₘM, but that doesn't seem to work.
  X = log(G, p, m) # X ∈ TₚM, # this one gives the correct hex.
  # X = log(G, m, p) # X ∈ TₘM, 
  return vee(LieAlgebra(G), X)
end

"""
    $SIGNATURES
Generic function that can be used in prior factors to calculate distance on Lie Groups. 
"""
function distancePrior(M::AbstractManifold, meas, p)
  return log(M, p, meas)
  # return distance(M, p, meas)
end

## ======================================================================================
## ManifoldFactor
## ======================================================================================

export ManifoldFactor
# DEV NOTES
# For now, `Z` is on the tangent space in coordinates at the point used in the factor.
# For groups just the lie algebra
# As transition it will be easier this way, we can reevaluate
struct ManifoldFactor{M <: AbstractManifold, T <: SamplableBelief} <: RelativeObservation
  M::M
  Z::T
end

DFG.getManifold(f::ManifoldFactor) = f.M

function getSample(cf::CalcFactor{<:ManifoldFactor{M, Z}}) where {M, Z}
  #TODO @assert dim == cf.factor.Z's dimension
  #TODO investigate use of SVector if small dims
  # if M isa ManifoldKernelDensity
  #   ret = sample(cf.factor.Z.belief)[1]
  # else
  #   ret = rand(cf.factor.Z)
  # end

  # ASSUME this function is only used for RelativeFactors which must use measurements as tangents
  ret = sampleTangent(cf.factor.M, cf.factor.Z)
  #return coordinates as we do not know the point here #TODO separate Lie group
  return ret
end

# function (cf::CalcFactor{<:ManifoldFactor{<:AbstractDecoratorManifold}})(Xc, p, q)
function (cf::CalcFactor{<:ManifoldFactor})(X, p, q)
  return measurement_residual(cf.factor.M, X, p, q)
end

## ======================================================================================
## adjoint factor - adjoint action applied to the measurement
## ======================================================================================
function Ad(::Union{typeof(SpecialEuclidean(2)), typeof(SpecialEuclidean(3))}, p, X)
  t = p.x[1]
  R = p.x[2]
  v = X.x[1]
  Ω = X.x[2]
  ArrayPartition(-R*Ω*R'*t + R*v, R*Ω*R')
end

function Ad(::typeof(SpecialEuclidean(3)), p)
  t = p.x[1]
  R = p.x[2]
  vcat(
      hcat(R, skew(t)*R),
      hcat(zero(SMatrix{3,3,Float64}), R)
  )
end

function Ad(::typeof(SpecialEuclidean(2)), p)
  t = p.x[1]
  R = p.x[2]
  vcat(
      hcat(R, -SA[0 -1; 1 0]*t),
      SA[0 0 1]
  )
end

function Ad(::Motion(2), p)
  t = p.x[1]
  R = p.x[2]
  vcat(
      hcat(R, -SA[0 -1; 1 0]*t),
      SA[0 0 1]
  )
end

struct AdFactor{F <: AbstractManifoldMinimize} <: AbstractManifoldMinimize
  factor::F
end

function (cf::CalcFactor{<:AdFactor})(Xϵ, p, q)
  # M = getManifold(cf.factor)
  # p,q ∈ M
  # Xϵ ∈ TϵM
  # ϵ = identity_element(M)
  # transform measurement from TϵM to TpM (global to local coordinates)
  # Adₚ⁻¹ = AdjointMatrix(M, p)⁻¹ = AdjointMatrix(M, p⁻¹)
  # Xp = Adₚ⁻¹ * Xϵᵛ
  # ad = Ad(M, inv(M, p))
  # Xp = Ad(M, inv(M, p), Xϵ)
  # Xp = adjoint_action(M, inv(M, p), Xϵ)
  #TODO is vector transport supposed to be the same?
  # Xp = vector_transport_to(M, ϵ, Xϵ, p)

  # Transform measurement covariance
  # ᵉΣₚ = Adₚ ᵖΣₚ Adₚᵀ
  #TODO test if transforming sqrt_iΣ is the same as Σ
  # Σ = ad * inv(cf.sqrt_iΣ^2) * ad'
  # sqrt_iΣ = convert(typeof(cf.sqrt_iΣ), sqrt(inv(Σ)))
  # sqrt_iΣ = convert(typeof(cf.sqrt_iΣ), ad * cf.sqrt_iΣ * ad')
  Xp = Xϵ

  child_cf = CalcFactorResidual(
    cf.faclbl,
    cf.factor.factor,
    cf.varOrder,
    cf.varOrderIdxs,
    cf.meas,
    cf.sqrt_iΣ,
    cf.cache,
  )
  return child_cf(Xp, p, q)
end

getMeasurementParametric(f::AdFactor) = getMeasurementParametric(f.factor)

getManifold(f::AdFactor) = getManifold(f.factor)
function getSample(cf::CalcFactor{<:AdFactor}) 
  M = getManifold(cf)
  return sampleTangent(M, cf.factor.factor.Z)
end


## ======================================================================================
## adjoint factor - adjoint action applied to the measurement
## ======================================================================================

# Adjoints defined in ApproxManifoldProducts
struct AdFactor{F <: RelativeObservation} <: RelativeObservation
  factor::F
end

function (cf::CalcFactor{<:AdFactor})(Xϵ, p, q)
  # M = getManifold(cf.factor)
  # p,q ∈ M
  # Xϵ ∈ TϵM
  # ϵ = identity_element(M)
  # transform measurement from TϵM to TpM (global to local coordinates)
  # Adₚ⁻¹ = AdjointMatrix(M, p)⁻¹ = AdjointMatrix(M, p⁻¹)
  # Xp = Adₚ⁻¹ * Xϵᵛ
  # ad = Ad(M, inv(M, p))
  # Xp = Ad(M, inv(M, p), Xϵ)
  # Xp = adjoint_action(M, inv(M, p), Xϵ)
  #TODO is vector transport supposed to be the same?
  # Xp = vector_transport_to(M, ϵ, Xϵ, p)

  # Transform measurement covariance
  # ᵉΣₚ = Adₚ ᵖΣₚ Adₚᵀ
  #TODO test if transforming sqrt_iΣ is the same as Σ
  # Σ = ad * inv(cf.sqrt_iΣ^2) * ad'
  # sqrt_iΣ = convert(typeof(cf.sqrt_iΣ), sqrt(inv(Σ)))
  # sqrt_iΣ = convert(typeof(cf.sqrt_iΣ), ad * cf.sqrt_iΣ * ad')
  Xp = Xϵ

  child_cf = CalcFactorResidual(
    cf.faclbl,
    cf.factor.factor,
    cf.varOrder,
    cf.varOrderIdxs,
    cf.meas,
    cf.sqrt_iΣ,
    cf.cache,
  )
  return child_cf(Xp, p, q)
end

getMeasurementParametric(f::AdFactor) = getMeasurementParametric(f.factor)

getManifold(f::AdFactor) = getManifold(f.factor)
function getSample(cf::CalcFactor{<:AdFactor})
  M = getManifold(cf)
  return sampleTangent(M, cf.factor.factor.Z)
end

## ======================================================================================
## ManifoldPrior
## ======================================================================================
export ManifoldPrior, PackedManifoldPrior

# `p` is a point on manifold `M`
# `Z` is a measurement at the tangent space of `p` on manifold `M` 
struct ManifoldPrior{M <: AbstractManifold, T <: SamplableBelief, P, B <: AbstractBasis} <:
       AbstractPriorObservation
  M::M
  p::P #NOTE This is a fixed point from where the measurement `Z` is made in coordinates on tangent TpM
  Z::T
  basis::B
  retract_method::AbstractRetractionMethod
end

function ManifoldPrior(M::AbstractLieGroup, p, Z)
  return ManifoldPrior(M, p, Z, DefaultLieAlgebraOrthogonalBasis(), MB.ExponentialRetraction())
end

DFG.getManifold(f::ManifoldPrior) = f.M

#TODO
# function ManifoldPrior(M::AbstractDecoratorManifold, Z::SamplableBelief)
#     # p = identity_element(M, #TOOD)
#     # similar to getPointIdentity(M)
#     return ManifoldPrior(M, Z, p)
# end

# ManifoldPrior{M}(Z::SamplableBelief, p) where M = ManifoldPrior{M, typeof(Z), typeof(p)}(Z, p)

function getSample(cf::CalcFactor{<:ManifoldPrior})
  Z = cf.factor.Z
  p = cf.factor.p
  M = cf.factor.M
  basis = cf.factor.basis
  retract_method = cf.factor.retract_method
  point = samplePoint(M, Z, p, basis, retract_method)

  return point
end

function getSample(cf::CalcFactor{<:ManifoldPrior{<:AbstractLieGroup}})
  Z = cf.factor.Z
  p = cf.factor.p
  M = cf.factor.M
  point = samplePoint(M, Z, p)

  return point
end

function getFactorMeasurementParametric(fac::ManifoldPrior)
  M = getManifold(fac)
  dims = manifold_dimension(M)
  meas = fac.p
  iΣ = convert(SMatrix{dims, dims}, invcov(fac.Z))
  return meas, iΣ
end

#TODO investigate SVector if small dims, this is slower
# dim = manifold_dimension(M)
# Xc = [SVector{dim}(rand(Z)) for _ in 1:N]

function (cf::CalcFactor{<:ManifoldPrior{<:AbstractLieGroup}})(m, p)
  M = cf.factor.M
  return prior_residual(M, m, p)
end

# dist²_Σ = ⟨X, Σ⁻¹*X'⟩
function mahalanobus_distance2(M, p, q, inv_Σ)
  Xc = log(M, p, q)
  return mahalanobus_distance2(Xc, inv_Σ)
end

function mahalanobus_distance2(M, X, inv_Σ)
  #TODO look to replace with inner(MM, p, X, inv_Σ*X)
  # Xc = get_coordinates(M, p, X, DefaultOrthogonalBasis())
  Xc = vee(M, p, X)
  return Xc' * inv_Σ * Xc
end

Base.@kwdef mutable struct PackedManifoldPrior <: AbstractPackedObservation
  varType::String
  p::Vector{Float64}  #NOTE This is a fixed point from where the measurement `Z` likely stored as a coordinate
  Z::PackedBelief
end

function convert(
  ::Union{Type{<:AbstractPackedObservation}, Type{<:PackedManifoldPrior}},
  obj::ManifoldPrior,
)
  #

  varT = typeModuleName(getStateKind(obj.M))

  c = AMP.makeCoordsFromPoint(obj.M, obj.p)

  # TODO convert all distributions to JSON
  Zst = convert(PackedBelief, obj.Z) # String

  return PackedManifoldPrior(varT, c, Zst)
end

function convert(
  ::Union{Type{<:AbstractObservation}, Type{<:ManifoldPrior}},
  obj::PackedManifoldPrior,
)
  #

  # piggy back on serialization of StateType rather than try serialize anything Manifolds.jl
  M = getTypeFromSerializationModule(obj.varType) |> getManifold

  # TODO this is too excessive
  e0 = getPointIdentity(M)
  # u0 = getPointIdentity(obj.varType)
  p = AMP.makePointFromCoords(M, obj.p, e0) #, u0)

  Z = convert(SamplableBelief, obj.Z)

  return ManifoldPrior(M, p, Z)
end

## ======================================================================================
## Generic Manifold Partial Prior
## ======================================================================================

function samplePointPartial(
  M::AbstractDecoratorManifold,
  z::Distribution,
  partial::Vector{Int},
  p = getPointIdentity(M),
  retraction_method::AbstractRetractionMethod = ExponentialRetraction(),
)
  dim = manifold_dimension(M)
  Xc = zeros(dim)
  Xc[partial] .= rand(z)
  X = hat(M, p, Xc)
  return retract(M, p, X, retraction_method)
end

struct ManifoldPriorPartial{M <: AbstractManifold, T <: SamplableBelief, P <: Tuple} <:
       AbstractPriorObservation
  M::M
  Z::T
  partial::P
end

DFG.getManifold(f::ManifoldPriorPartial) = f.M

function getSample(cf::CalcFactor{<:ManifoldPriorPartial})
  Z = cf.factor.Z
  M = getManifold(cf)
  partial = collect(cf.factor.partial)

  return (samplePointPartial(M, Z, partial),)
end