from typing import Annotated, Literal

from typing_extensions import TypeAliasType

from pydantic import BaseModel, Field, model_validator

from olmoearth_run.shared.split_proportions import validate_split_proportions


class RandomDataSplitter(BaseModel):
    """Data splitter that assigns splits based on random sampling."""
    name: Literal["random"]
    train_prop: float = Field(ge=0, le=1, description="The proportion of data to assign to the training split.")
    val_prop: float = Field(ge=0, le=1, description="The proportion of data to assign to the validation split.")
    test_prop: float = Field(ge=0, le=1, description="The proportion of data to assign to the test split.")
    seed: int | None = Field(default=42, description="The seed for the random number generator.")

    @model_validator(mode='after')
    def validate_proportions_sum_to_one(self) -> 'RandomDataSplitter':
        """Ensure that train_prop + val_prop + test_prop == 1"""
        validate_split_proportions(self.train_prop, self.val_prop, self.test_prop)
        return self


class SpatialDataSplitter(BaseModel):
    """Data splitter that assigns splits based on spatial grid cell location."""
    name: Literal["spatial"]
    grid_size: int = Field(ge=1, description="The size of the grid cells in pixels.")
    train_prop: float = Field(ge=0, le=1, description="The proportion of data to assign to the training split.")
    val_prop: float = Field(ge=0, le=1, description="The proportion of data to assign to the validation split.")
    test_prop: float = Field(ge=0, le=1, description="The proportion of data to assign to the test split.")

    @model_validator(mode='after')
    def validate_proportions_sum_to_one(self) -> 'SpatialDataSplitter':
        """Ensure that train_prop + val_prop + test_prop == 1."""
        validate_split_proportions(self.train_prop, self.val_prop, self.test_prop)
        return self


DataSplitter = TypeAliasType(
    "DataSplitter",
    Annotated[
        RandomDataSplitter | SpatialDataSplitter,
        Field(discriminator="name")
    ]
)
