"""multistate joint modeling package.

This package provides tools for multistate joint modeling with PyTorch.
"""

from .functions import (
    Exponential,
    Gompertz,
    LogNormal,
    Net,
    Weibull,
    clock_forward,
    clock_reset,
    gamma_plus_b,
    gamma_x_plus_b,
    identity,
    linear,
)
from .jobs import (
    AdamL1Proximal,
    ComputeCriteria,
    ComputeEBEs,
    ComputeFIM,
    Fit,
    GradStop,
    LogParamsHistory,
    MCMCDiagnostics,
    ParamStop,
    PredictSurvLogps,
    PredictTrajectories,
    PredictY,
    Scheduling,
    SwitchParams,
    ValueStop,
)
from .model._base import MultiStateJointModel
from .typedefs import (
    BaseHazardFn,
    BucketData,
    ClockMethod,
    IndividualEffectsFn,
    Info,
    Job,
    LinkFn,
    MatRepr,
    Metrics,
    ModelData,
    ModelDesign,
    ModelParams,
    RegressionFn,
    SampleData,
)
from .utils import build_buckets, cov_from_repr, get_dtype, repr_from_cov, set_dtype

__author__ = "Félix Laplante"
__email__ = "felixlaplante0@gmail.com"
__license__ = "MIT"

__all__ = [
    "AdamL1Proximal",
    "BaseHazardFn",
    "BucketData",
    "ClockMethod",
    "ComputeCriteria",
    "ComputeEBEs",
    "ComputeFIM",
    "Exponential",
    "Fit",
    "Gompertz",
    "GradStop",
    "IndividualEffectsFn",
    "Info",
    "Job",
    "LinkFn",
    "LogNormal",
    "LogParamsHistory",
    "MCMCDiagnostics",
    "MatRepr",
    "Metrics",
    "ModelData",
    "ModelDesign",
    "ModelParams",
    "MultiStateJointModel",
    "Net",
    "ParamStop",
    "PredictSurvLogps",
    "PredictTrajectories",
    "PredictY",
    "RegressionFn",
    "SampleData",
    "Scheduling",
    "SwitchParams",
    "ValueStop",
    "Weibull",
    "build_buckets",
    "clock_forward",
    "clock_reset",
    "cov_from_repr",
    "gamma_plus_b",
    "gamma_x_plus_b",
    "get_dtype",
    "identity",
    "linear",
    "repr_from_cov",
    "set_dtype",
]
