"""Hackable surface for the multiscale polynomial memory research code."""
from __future__ import annotations

from .integration import IntegrationResult, integrate_system
from .models import AssocMemHiPPO, SalienceHiPPO, VanillaHiPPO, legendre_orthonormal_basis01
from .polynomials import (
    EvalFn,
    SystemParams,
    fourier_val,
    get_L_vector,
    get_output_vector,
    get_system_params,
    legval,
    legsval,
    legval_jax,
    legsval_jax,
    make_m_mat,
)
from .synthetic import (
    get_ou_sample,
    oracle_mse_censored_task,
    oracle_info_gain,
    oracle_mse_rbf_periodic_censor,
    oracle_info_gain_rbf,
    make_copying_task,
    make_censored_task,
    make_induction_head_multioutput_s2s_task,
    make_induction_head_seq_to_seq_task,
    make_induction_head_task,
    make_implicit_measure_task,
    make_multiplexing_task,
    whitesignal,
    wray_and_green_output,
    make_simple_repetition_task
)
from .utils import time_block, unvec, vec

__all__ = [
    "AssocMemHiPPO",
    "EvalFn",
    "IntegrationResult",
    "SystemParams",
    "fourier_val",
    "get_L_vector",
    "get_output_vector",
    "get_ou_sample",
    "get_system_params",
    "SalienceHiPPO",
    "VanillaHiPPO",
    "integrate_system",
    "legendre_orthonormal_basis01",
    "legval",
    "legsval",
    "make_induction_head_multioutput_s2s_task",
    "make_induction_head_seq_to_seq_task",
    "make_induction_head_task",
    "make_implicit_measure_task",
    "make_m_mat",
    "make_multiplexing_task",
    "time_block",
    "unvec",
    "vec",
]