api/schema/input_data.py

from typing import Literal, Optional

from pydantic import BaseModel, Field, StrictBool, model_validator, StrictFloat

from api.schema.input_helpers import (
    JOB_CATEGORY_MODEL,
    age_category_map,
    education_map,
    map_customer_job,
)


class Education(BaseModel):
    type: Literal["school", "university", "course", "illiterate"]
    level: Optional[Literal["primary", "middle", "high"]] = None
    grade: Optional[int] = None

    @model_validator(mode="before")
    def validate_education(cls, values):
        type_ = values.get("type")
        level = values.get("level")
        grade = values.get("grade")

        if type_ == "school":
            if level is None or grade is None:
                raise ValueError(
                    "For type 'school', both level and grade must be provided"
                )
        else:
            if level is not None or grade is not None:
                raise ValueError(
                    f"For type '{type_}', level and grade must not be provided"
                )
        return values


class PersonalInfo(BaseModel):
    age: int = Field(..., ge=0, le=120)
    job: Literal[
        "blue_collar",
        "housemaid",
        "services",
        "admin",
        "technician",
        "management",
        "self_employed",
        "entrepreneur",
        "unemployed",
        "student",
    ]
    marital: Literal["married", "single"]
    education: Education

    # Derived fields
    job_category: Optional[Literal["cat1", "cat2", "cat3", "cat4", "other"]] = None
    age_category: Optional[
        Literal[
            "struggling",
            "stable",
            "about to retire",
            "old age",
            "counting a last breathe",
        ]
    ] = None

    education_str: Optional[str] = None

    @model_validator(mode="before")
    def derive_categories(cls, values):
        job_input = values.get("job")
        age_input = values.get("age")
        edu_dict = values.get("education")

        job_model_name = map_customer_job(job_input)
        job_cat = JOB_CATEGORY_MODEL.get(job_model_name, "other")
        age_cat = age_category_map(age_input)

        edu_type = edu_dict.get("type") if edu_dict else None
        edu_level = edu_dict.get("level") if edu_dict else None
        edu_grade = edu_dict.get("grade") if edu_dict else None

        values["education_str"] = education_map(
            type=edu_type, level=edu_level, grade=edu_grade
        )

        values["job_category"] = job_cat
        values["age_category"] = age_cat
        return values


class FinancialInfo(BaseModel):
    default: StrictBool
    housing: StrictBool
    loan: StrictBool

    @model_validator(mode="after")
    def transform_bool_to_int(cls, values):
        values.default = int(values.default)
        values.housing = int(values.housing)
        values.loan = int(values.loan)
        return values


class ContactInfo(BaseModel):
    contact: Literal["cellular", "telephone"]
    day_of_week: Literal["mon", "tue", "wed", "thu", "fri"]
    month: Literal[
        "mar",
        "apr",
        "may",
        "jun",
        "jul",
        "aug",
        "sep",
        "oct",
        "nov",
        "dec",
    ]


class CampaignInfo(BaseModel):
    campaign: int = Field(..., ge=0)
    previous: int = Field(..., ge=0)
    poutcome: Literal["success", "failure", "nonexistent"]
    

class MacroInfo(BaseModel):
    emp_var_rate: StrictFloat = Field(..., alias="employment_variation_rate")
    euribor3m: StrictFloat = Field(..., alias="euribor_3m_rate")
    nr_employed: StrictFloat = Field(..., alias="number_employed")
    cons_price_idx: StrictFloat = Field(..., alias="consumer_price_index")
    cons_conf_idx: float = Field(..., alias="consumer_confidence_index")

class InputData(BaseModel):
    personal_info: PersonalInfo
    financial_info: FinancialInfo
    contact_info: ContactInfo
    campaign_info: CampaignInfo
    macro_info: MacroInfo