src/data/schema.py

"""Data schemas for structural mechanics dataset.

Pydantic models enforce type safety and validation across the entire pipeline.
All quantities use SI units (meters, Pascals, Newtons).
"""

from enum import Enum
from typing import Optional

from pydantic import BaseModel, Field, field_validator


class ProblemFamily(str, Enum):
    BEAM = "beam"
    PLATE = "plate"
    VESSEL = "vessel"


class BeamConfig(str, Enum):
    SS_POINT = "beam_ss_point"
    SS_UDL = "beam_ss_udl"
    CANTILEVER_POINT = "beam_cantilever_point"
    CANTILEVER_UDL = "beam_cantilever_udl"
    FIXED_POINT = "beam_fixed_point"
    FIXED_UDL = "beam_fixed_udl"


class PlateConfig(str, Enum):
    SS_UNIFORM = "plate_ss_uniform"
    FIXED_UNIFORM = "plate_fixed_uniform"


class VesselConfig(str, Enum):
    CYLINDER = "vessel_cylinder"
    SPHERE = "vessel_sphere"


class SafetyCategory(str, Enum):
    SAFE = "safe"          # SF >= 2.0
    MARGINAL = "marginal"  # 1.0 <= SF < 2.0
    FAILURE = "failure"    # SF < 1.0


class MaterialProperties(BaseModel):
    """Material properties in SI units."""

    elastic_modulus: float = Field(..., gt=0, description="Young's modulus E [Pa]")
    poisson_ratio: float = Field(..., gt=0, lt=0.5, description="Poisson's ratio nu [-]")
    yield_strength: float = Field(..., gt=0, description="Yield strength sigma_y [Pa]")
    density: float = Field(..., gt=0, description="Density rho [kg/m^3]")


class BeamParameters(BaseModel):
    """Input parameters for Euler-Bernoulli beam analysis."""

    length: float = Field(..., gt=0, description="Beam span L [m]")
    width: float = Field(..., gt=0, description="Cross-section width b [m]")
    height: float = Field(..., gt=0, description="Cross-section height h [m]")
    material: MaterialProperties
    point_load: float = Field(default=0.0, ge=0, description="Point load P [N]")
    distributed_load: float = Field(default=0.0, ge=0, description="Distributed load w [N/m]")

    @property
    def moment_of_inertia(self) -> float:
        """Second moment of area I = bh^3/12 [m^4]."""
        return self.width * self.height**3 / 12.0

    @property
    def section_modulus(self) -> float:
        """Section modulus S = bh^2/6 [m^3]."""
        return self.width * self.height**2 / 6.0

    @property
    def cross_section_area(self) -> float:
        """Cross-section area A = bh [m^2]."""
        return self.width * self.height


class PlateParameters(BaseModel):
    """Input parameters for Kirchhoff-Love plate analysis."""

    length_a: float = Field(..., gt=0, description="Plate dimension a [m]")
    length_b: float = Field(..., gt=0, description="Plate dimension b [m]")
    thickness: float = Field(..., gt=0, description="Plate thickness t [m]")
    material: MaterialProperties
    pressure: float = Field(..., gt=0, description="Uniform pressure q [Pa]")

    @property
    def flexural_rigidity(self) -> float:
        """Flexural rigidity D = Et^3 / (12(1 - nu^2)) [N*m]."""
        E = self.material.elastic_modulus
        nu = self.material.poisson_ratio
        t = self.thickness
        return E * t**3 / (12.0 * (1.0 - nu**2))

    @property
    def aspect_ratio(self) -> float:
        """Aspect ratio a/b [-]."""
        return self.length_a / self.length_b


class VesselParameters(BaseModel):
    """Input parameters for thick-walled pressure vessel analysis (Lame equations)."""

    inner_radius: float = Field(..., gt=0, description="Inner radius r_i [m]")
    outer_radius: float = Field(..., gt=0, description="Outer radius r_o [m]")
    length: Optional[float] = Field(default=None, gt=0, description="Cylinder length L [m]")
    material: MaterialProperties
    internal_pressure: float = Field(..., gt=0, description="Internal pressure p [Pa]")

    @field_validator("outer_radius")
    @classmethod
    def outer_must_exceed_inner(cls, v: float, info: dict) -> float:
        if "inner_radius" in info.data and v <= info.data["inner_radius"]:
            raise ValueError("outer_radius must be greater than inner_radius")
        return v

    @property
    def radius_ratio(self) -> float:
        """Radius ratio r_o/r_i [-]."""
        return self.outer_radius / self.inner_radius

    @property
    def wall_thickness(self) -> float:
        """Wall thickness t = r_o - r_i [m]."""
        return self.outer_radius - self.inner_radius


class SolutionResult(BaseModel):
    """Output from an analytical solver."""

    max_stress: float = Field(..., ge=0, description="Maximum stress [Pa]")
    max_deflection: float = Field(..., ge=0, description="Maximum deflection [m]")
    safety_factor: float = Field(..., ge=0, description="Safety factor = sigma_y / max_stress [-]")
    safety_category: SafetyCategory

    @classmethod
    def from_stress(cls, max_stress: float, max_deflection: float, yield_strength: float) -> "SolutionResult":
        """Construct result and auto-classify safety category."""
        sf = yield_strength / max_stress if max_stress > 0 else float("inf")

        if sf >= 2.0:
            category = SafetyCategory.SAFE
        elif sf >= 1.0:
            category = SafetyCategory.MARGINAL
        else:
            category = SafetyCategory.FAILURE

        return cls(
            max_stress=max_stress,
            max_deflection=abs(max_deflection),
            safety_factor=sf,
            safety_category=category,
        )


class DatasetSample(BaseModel):
    """A single row in the dataset. Flat structure for Parquet compatibility."""

    sample_id: str
    problem_family: ProblemFamily
    config_id: str

    # Geometry (SI: meters)
    length: float
    width: Optional[float] = None
    height: Optional[float] = None
    inner_radius: Optional[float] = None
    outer_radius: Optional[float] = None
    thickness: Optional[float] = None   # plate thickness

    # Material (SI)
    elastic_modulus: float
    poisson_ratio: float
    yield_strength: float
    density: float

    # Loading (SI)
    point_load: float = 0.0
    distributed_load: float = 0.0
    internal_pressure: float = 0.0
    pressure: float = 0.0  # plate uniform pressure

    # Derived geometry
    moment_of_inertia: Optional[float] = None
    section_modulus: Optional[float] = None
    cross_section_area: Optional[float] = None

    # Outputs
    max_stress: float
    max_deflection: float
    safety_factor: float
    safety_category: SafetyCategory