from typing import Annotated, Literal

from typing_extensions import TypeAliasType

from pydantic import BaseModel, Field

from olmoearth_run.shared.models.config.model.decoders import SegmentationDecoder


class SegmentationTask(BaseModel):
    name: Literal["segmentation"]
    decoder: SegmentationDecoder


# TODO: classification, per-pixel regression, regression, and detection tasks


ModelTask = TypeAliasType(
    "ModelTask",
    Annotated[
        SegmentationTask,
        Field(discriminator="name")
    ]
)
