"""
Utilities module: provides the public API for utility functions used by core algorithms.
Internal helper functions should be imported directly from their submodules.

Organized into focused modules:
- offline_helpers: Utilities specific to offline learning
- online_helpers: Utilities specific to online learning
- misc_helpers: any helpers not strictly online or offline. e.g. trajectory manipulation
- embeddings: Embedding functions for trajectories
"""

from .offline_helpers import (
    calc_offlineradius,
    generate_confidence_set_deterministic_via_noise_matrices,
    generate_confidence_set_deterministic_via_rejection_sampling,
    calc_d_pi_BC,
    generate_offline_trajectories,
)

from .online_helpers import (
    precompute_phi_Bhats,
    calc_gamma_t,
    calc_empirical_counts,
    get_policy_pair_that_maximizes_uncertainty,
    generate_policy_pair_rollouts,
    annotate_buffer,
    learn_w_MLE,
    project_w,
    find_most_preferred_policy,
    calc_regret,
    loop_iteration_logging,
    initial_loop_earlystop_logging,
    final_iteration_logging,
    compute_rewards_traj,
    calc_avg_reward,
)

from .misc_helpers import (
    rollout_policy_in_env,
    pad_metrics,
    get_fig_names,
    dirs_and_loads,
    short_env_name,
    save_metrics,
)

from .embeddings import create_embeddings

__all__ = [
    # Offline learning public API (needed in offline_learning.py)
    "calc_offlineradius",
    "generate_confidence_set_deterministic_via_noise_matrices",
    "generate_confidence_set_deterministic_via_rejection_sampling",
    "calc_d_pi_BC",
    "generate_offline_trajectories",
    # Online learning public API (needed in online_learning.py)
    "precompute_phi_Bhats",
    "calc_gamma_t",
    "calc_empirical_counts",
    "get_policy_pair_that_maximizes_uncertainty",
    "generate_policy_pair_rollouts",
    "annotate_buffer",
    "learn_w_MLE",
    "project_w",
    "find_most_preferred_policy",
    "calc_regret",
    "loop_iteration_logging",
    "initial_loop_earlystop_logging",
    "final_iteration_logging",
    "compute_rewards_traj",
    "calc_avg_reward",
    # Misc helpers
    "rollout_policy_in_env",
    "pad_metrics",
    "get_fig_names",
    "dirs_and_loads",
    "short_env_name",
    "save_metrics",
    # Embeddings
    "create_embeddings",
]
