from typing import List, Optional, Dict

from pydantic import Field, BaseModel

from magdi_data.dataset.nifti_occ_inst_dataset_meta import NiftiOccInstDatasetMeta


class DescriptiveMetaData(BaseModel):
    latin_name: str = Field()
    line_name: str = Field()
    structure: str = Field()
    dap: str = Field()
    device: str = Field()
    coil: str = Field()
    measurement_channel: str = Field()


class PlantNiftiOccInstDatasetMeta(NiftiOccInstDatasetMeta):

    descriptive_metadata: DescriptiveMetaData = Field(
        default_factory=DescriptiveMetaData  # type: ignore
    )
    tags: List[str] = Field(default=["MRI", "3D", "Image"])
    task_categories: List[str] = Field(default=["image-segmentation"])
    labels: Optional[Dict[str, str]] = Field(
        default={"0": "background", "1": "embryo", "2": "endosperm", "3": "aleuron"}
    )
    image_format: str = Field(default=".nii.gz")
    annotation_format: Optional[str] = Field(default=".nii.gz")
    image_data_type: str = Field(default="int16")
    annotation_data_type: Optional[str] = Field(default="uint8")
    range: List[int] = Field(default=[0, 16000])
    resolution_voxel_size: List[float] = Field(default=[0.1, 0.1, 0.1])
