"""Definition of ModelData data model."""

from __future__ import annotations

from .angle import AngleData
from .contact import ContactData
from .data_model import DataModel
from .distance import DistanceData
from .framing import FramingData
from .information_value import InformationValueData
from .modality_color import ModalityColorData
from .modality_depth import ModalityDepthData
from .modality_lighting import ModalityLightingData
from .point_of_view import PointOfViewData
from .salience import SalienceData
from .visual_syntax import VisualSyntaxData


class ModelData(DataModel):
    """ModelData model for data IO with combined ML model."""

    visual_syntax: VisualSyntaxData
    contact: ContactData
    angle: AngleData
    point_of_view: PointOfViewData
    distance: DistanceData
    modality_lighting: ModalityLightingData
    modality_color: ModalityColorData
    modality_depth: ModalityDepthData
    information_value: InformationValueData
    framing: FramingData
    salience: SalienceData

    @classmethod
    def from_random(cls) -> ModelData:
        """Instantiate with random numbers."""
        kwargs = {
            field: field_info.annotation.from_random()  # type: ignore
            for field, field_info in cls.model_fields.items()
        }
        return cls(**kwargs)

    @classmethod
    def from_prediction_dict(
        cls,
        prediction_dict: dict,
    ) -> ModelData:
        """Instantiate from prediction dictionary."""
        kwargs = {
            'visual_syntax': VisualSyntaxData.from_tensor(
                prediction_dict['visual_syntax'],
            ),
            'contact': ContactData.from_tensor(
                prediction_dict['contact'],
            ),
            'angle': AngleData.from_tensor(
                prediction_dict['angle'],
            ),
            'point_of_view': PointOfViewData.from_tensor(
                prediction_dict['point_of_view'],
            ),
            'distance': DistanceData.from_tensor(
                prediction_dict['distance'],
            ),
            'modality_lighting': ModalityLightingData.from_tensor(
                prediction_dict['modality_lighting'],
            ),
            'modality_color': ModalityColorData.from_tensor(
                prediction_dict['modality_color'],
            ),
            'modality_depth': ModalityDepthData.from_tensor(
                prediction_dict['modality_depth'],
            ),
            'information_value': InformationValueData.from_tensor(
                prediction_dict['information_value'],
            ),
            'framing': FramingData.from_tensor(
                prediction_dict['framing'],
            ),
            'salience': SalienceData.from_tensor(
                prediction_dict['salience'],
            ),
        }
        return cls(**kwargs)

    @classmethod
    def from_annotations(
        cls,
        visual_syntax: str,
        contact: str,
        angle: str,
        point_of_view: str,
        distance: str,
        modality_lighting: str,
        modality_color: str,
        modality_depth: str,
        information_value: str,
        framing: str,
        salience: str,
    ) -> ModelData:
        """Instantiate from annotation."""
        kwargs = {
            'visual_syntax': VisualSyntaxData.from_choice(visual_syntax),
            'contact': ContactData.from_choice(contact),
            'angle': AngleData.from_choice(angle),
            'point_of_view': PointOfViewData.from_choice(point_of_view),
            'distance': DistanceData.from_choice(distance),
            'modality_lighting': ModalityLightingData.from_choice(modality_lighting),
            'modality_color': ModalityColorData.from_choice(modality_color),
            'modality_depth': ModalityDepthData.from_choice(modality_depth),
            'information_value': InformationValueData.from_choice(information_value),
            'framing': FramingData.from_choice(framing),
            'salience': SalienceData.from_choice(salience),
        }
        return cls(**kwargs)