"""Objective selection shared by training and the standalone inference bundle."""
from trace_imf import TraceIMF
from multi_trace_imf import MultiTraceIMF


def build_objective(cfg):
    common = dict(channels=cfg["model"]["in_channels"],
                  image_size=cfg["model"]["input_size"],
                  num_classes=cfg["model"]["num_classes"])
    name = cfg.get("objective", "trace_imf")
    if name == "multi_trace_imf":
        return MultiTraceIMF(**common, **cfg["multi_trace_imf"])
    if name == "trace_imf":
        return TraceIMF(**common, **cfg["trace_imf"])
    raise ValueError(f"Unknown objective: {name}")