HybridValues

import gtsam
import numpy as np

from gtsam import HybridValues, VectorValues, DiscreteValues, Values, Pose2
from gtsam.symbol_shorthand import V, D, M  # Use V, D, M for keys

# 1. Default constructor (empty)
hybrid_values_empty = HybridValues()
print("Empty HybridValues:")
hybrid_values_empty.print()

# 2. From VectorValues and DiscreteValues
vec_vals = VectorValues()
vec_vals.insert(V(0), np.array([1.0, 2.0]))
vec_vals.insert(V(1), np.array([3.0]))

disc_vals = DiscreteValues()
disc_vals[D(0)] = 1
disc_vals[D(1)] = 0

hybrid_values_vd = HybridValues(vec_vals, disc_vals)
print("\nHybridValues from VectorValues and DiscreteValues:")
hybrid_values_vd.print()

# 3. From VectorValues, DiscreteValues, and Nonlinear Values
nonlinear_vals = Values()
nonlinear_vals.insert(M(5), Pose2(1, 2, 0.3)) # Example nonlinear type

hybrid_values_all = HybridValues(vec_vals, disc_vals, nonlinear_vals)
print("\nHybridValues from all three types:")
hybrid_values_all.print()

# Access underlying containers
cont_vals = hybrid_values_all.continuous()
disc_vals_acc = hybrid_values_all.discrete()
nonlin_vals_acc = hybrid_values_all.nonlinear()

print(f"\nAccessed Continuous Values size: {cont_vals.size()}")
cont_vals.print("Accessed Continuous Values:")
print(f"Accessed Discrete Values size: {len(disc_vals_acc)}") # DiscreteValues acts like dict
gtsam.PrintDiscreteValues(disc_vals_acc,"Accessed Discrete Values:")
print(f"Accessed Nonlinear Values size: {nonlin_vals_acc.size()}")
nonlin_vals_acc.print("Accessed Nonlinear Values:")

# Check existence
print(f"\nExists Vector V(0)? {hybrid_values_all.existsVector(V(0))}")
print(f"Exists Discrete D(1)? {hybrid_values_all.existsDiscrete(D(1))}")
print(f"Exists Nonlinear M(5)? {hybrid_values_all.existsNonlinear(M(5))}")
print(f"Exists Vector D(0)? {hybrid_values_all.existsVector(D(0))}") # False, D(0) is a discrete key

# exists() checks across all types (nonlinear, then vector, then discrete)
print(f"Exists V(0) (any type)? {hybrid_values_all.exists(V(0))}") # Checks VectorValues
print(f"Exists D(0) (any type)? {hybrid_values_all.exists(D(0))}") # Checks DiscreteValues
print(f"Exists M(5) (any type)? {hybrid_values_all.exists(M(5))}") # Checks Nonlinear Values

# Access specific values
print(f"\nValue at V(0): {hybrid_values_all.at(V(0))}")
print(f"Value at D(0) (via atDiscrete): {hybrid_values_all.atDiscrete(D(0))}") 
print(f"Value at D(0) (via discrete() dict access): {hybrid_values_all.discrete()[D(0)]}")
print(f"Value at M(5): {hybrid_values_all.nonlinear().atPose2(M(5))}") # Use type-specific getter from nonlinear() part

hv = HybridValues()

# Insert individual values
hv.insert(V(10), np.array([1.0, 2.0])) # Vector value, goes into hv.continuous()
hv.insert(D(10), 1)                    # Discrete value, goes into hv.discrete()
hv.insertNonlinear(M(11), Pose2(0.1, 0.2, 0.03)) # Nonlinear value, goes into hv.nonlinear()

print("After individual inserts:")
hv.print()

# Insert from containers
new_vec = VectorValues()
new_vec.insert(V(12), np.array([5.0]))
new_disc = DiscreteValues()
new_disc[D(11)] = 0
new_nonlin = Values()
new_nonlin.insert(M(15), Pose2(1,1,0))

hv.insert(new_vec)      # Merges into hv.continuous()
hv.insert(new_disc)     # Merges into hv.discrete()
hv.insert(new_nonlin)   # Merges into hv.nonlinear()
print("\nAfter container inserts:")
hv.print()

# Update existing values
update_vec = VectorValues()
update_vec.insert(V(10), np.array([99.0, 98.0]))
update_disc = DiscreteValues()
update_disc[D(10)] = 2
update_nonlin = Values()
update_nonlin.insert(M(11), Pose2(0.5,0.6,0.07))

hv.update(update_vec)
hv.update(update_disc)
hv.update(update_nonlin)
print("\nAfter update:")
hv.print()

# Retract (applies delta to the nonlinearValues_ part of HybridValues)
# Note: The continuous_ (VectorValues) part is NOT retracted by HybridValues.retract itself.
delta_nl = VectorValues() # Deltas are always VectorValues
# Create a delta for the Pose2 at M(11). M(11) current value: Pose2(0.5,0.6,0.07)
delta_pose2_M11 = np.array([0.05, -0.05, 0.01]) # dx, dy, dtheta
delta_nl.insert(M(11), delta_pose2_M11)

hv_retracted = hv.retract(delta_nl)
print("\nAfter retract (M(11) should change, V(10) should not):")
hv_retracted.print()

# Insert or assign
# Replaces if exists, inserts if not.
hv.insert_or_assign(V(10), np.array([100.0, 101.0])) # Overwrites V(10) in continuous_
hv.insert_or_assign(D(12), 1) # Inserts D(12) in discrete_
hv.insert_or_assign(V(13), np.array([13.0])) # Inserts V(13) in continuous_
# Note: insert_or_assign for nonlinear types is not directly on HybridValues.
# You would typically do this on the underlying Values container:
# hv.nonlinear().insert_or_assign(M(11), Pose2(...)) # if Values had insert_or_assign
# Or, more commonly, for nonlinear: hv.nonlinear().update(M(11), Pose2(...))
print("\nAfter insert_or_assign (for V and D keys):")
hv.print()