mappingservice/models.py

import logging.config
from enum import Enum
from typing import List, Union

from pydantic import BaseModel, validator

logger = logging.getLogger(__name__)


class PredictionResponse(BaseModel):
    label: str
    score: float


class BedType(str, Enum):
    none = ""
    single = "single"
    double = "double"
    queen = "queen"
    king = "king"
    california_king = "california_king"
    bunk = "bunk"
    sofa = "sofa"
    rollaway = "rollaway"
    futon = "futon"


class BedData(BaseModel):
    type: Union[BedType, None] = None
    count: Union[int, None] = None

    @validator("count", pre=True, always=True)
    def validate_count(cls, v):
        return v if v and v > 0 else None

    @validator("type", pre=True, always=True)
    def validate_type(cls, v):
        try:
            return BedType(v.lower()) if v else None
        except ValueError:
            return None

    @classmethod
    def create_valid_bed(cls, type: Union[str, None], count: Union[int, None]):
        if type and count and count > 0:
            try:
                type_enum = BedType(type.lower())
                return cls(type=type_enum, count=count)
            except ValueError:
                logger.error(f"Invalid bed type: {type}")
        return None

    def __init__(self, **data):
        variations = {
            "individual": BedType.single,
            "camaindividual": BedType.single,
            "doble": BedType.double,
            "camadoble": BedType.double,
            "reina": BedType.queen,
            "queenbed": BedType.queen,
            "rey": BedType.king,
            "californiaking": BedType.california_king,
            "litera": BedType.bunk,
            "sofá": BedType.sofa,
            "sofacama": BedType.sofa,
            "plegable": BedType.rollaway,
            "futón": BedType.futon,
            "twin": BedType.single,
            "twinbed": BedType.single,
            "singlebed": BedType.single,
            "doublebed": BedType.double,
            "largedouble": BedType.queen,
            "extralargedouble": BedType.king,
            "bunkbed": BedType.bunk,
            "couch": BedType.sofa,
            "airmattress": BedType.futon,
            "floormattress": BedType.futon,
        }

        bed_type = data.get("type")
        if bed_type:
            normalized_bed_type = bed_type.replace(" ", "").lower()
            data["type"] = variations.get(normalized_bed_type, bed_type)
        super().__init__(**data)


class RoomData(BaseModel):
    room_description: str
    beds: Union[BedData, None]


class Predictions(BaseModel):
    type: PredictionResponse
    category: PredictionResponse
    environment: List[PredictionResponse]
    feature: List[PredictionResponse]
    view: List[PredictionResponse]
    language_detected: str
    beds: List[BedData]



class AllPredictionsResponse(BaseModel):
    predictions: Predictions