from dataclasses import dataclass

from olmoearth_pretrain.data.constants import Modality as OlmoEarthPretrainModality  # type: ignore[import-untyped]


@dataclass(frozen=True)
class ModalitiesSpec:
    """Maps our modality types (e.g. Sentinel1RTC, Sentinel2L2A) to encoder-specific values."""

    sentinel1_rtc_name: str
    sentinel1_rtc_band_order: list[str]
    sentinel2_l2a_name: str
    sentinel2_l2a_band_order: list[str]


# Authoritative modality values from olmoearth_pretrain v1
OLMOEARTH_MODALITIES_SPEC = ModalitiesSpec(
    sentinel1_rtc_name=OlmoEarthPretrainModality.SENTINEL1.name,
    sentinel1_rtc_band_order=list(OlmoEarthPretrainModality.SENTINEL1.band_order),
    sentinel2_l2a_name=OlmoEarthPretrainModality.SENTINEL2_L2A.name,
    sentinel2_l2a_band_order=list(OlmoEarthPretrainModality.SENTINEL2_L2A.band_order),
)

# rslearn's task classes (SegmentationTask, ClassificationTask, etc.) all hardcode
# raw_inputs["targets"] in their process_inputs methods. When using SingleTaskModel,
# the input key must be "targets". When using MultiTask, input_mapping remaps to "targets".
TARGET = "targets"

# rslearn transforms use "target" as a special prefix for image_selectors,
# independent of the input key name. Format: "target/{field}" for SingleTask.
TARGET_SELECTOR_PREFIX = "target"

# STAC sar:polarizations query value (uppercase, unlike bands)
STAC_SAR_POLARIZATIONS = ["VV", "VH"]

# OlmoEarth models are trained on Sentinel-2 data at 10m/pixel resolution
OLMOEARTH_ENCODER_RESOLUTION_METERS = 10.0
OLMOEARTH_MODULE_SELECTOR: list[str | int] = ["model", "encoder", 0]

LABEL_LAYER_NAME = "labels"
OUTPUT_LAYER_NAME = "output"

STAC_EO_CLOUD_COVER = "eo:cloud_cover"
STAC_SAT_ORBIT_STATE_ASCENDING = "ascending"
STAC_SAT_ORBIT_STATE_DESCENDING = "descending"

DATA_SPLIT_KEY = "split"
DEFAULT_GROUP_NAME = "default"
