Blame for src/config/schema.py

"""Pydantic configuration schema for the synthetic document generation pipeline.

Every user-facing parameter is validated here. Numeric placement/augmentation
parameters accept either a fixed scalar or a ``{min, max}`` range (resolved
per-sample by a seeded RNG, see :mod:`src.utils.random_utils`).

The individual augmentation effects requested by the product spec are
consolidated into configurable *families* where an effect is a strict
specialization of another (paper stains + water stains -> ``stain`` with a
``type`` selector; edge/corner/fold shadows -> ``shadow`` with a ``type``
selector; crumpling + wrinkles -> ``crumple``; general discoloration + paper
yellowing -> ``discoloration`` with a ``tone`` selector). This keeps the
config surface coherent while still reaching every requested visual effect.
"""

from __future__ import annotations

from typing import Literal, Union

from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator

NumericOrRange_T = Union[float, "Range"]


class Range(BaseModel):
    """An inclusive ``[min, max]`` range a value is uniformly sampled from."""

    model_config = ConfigDict(extra="forbid")

    min: float
    max: float

    @model_validator(mode="after")
    def _check_order(self) -> "Range":
        if self.max < self.min:
            raise ValueError(f"Range max ({self.max}) must be >= min ({self.min})")
        return self


NumericOrRange = Union[float, Range]


class AugmentationParam(BaseModel):
    """Base schema shared by every augmentation effect.

    Extra keys are allowed so effect-specific parameters (e.g. ``type`` for
    :class:`ShadowParam`, ``tone`` for discoloration) can be declared inline
    in YAML without needing a bespoke subclass for every single effect.
    """

    model_config = ConfigDict(extra="allow")

    enabled: bool = False
    strength: NumericOrRange = 0.0
    probability: float = Field(default=1.0, ge=0.0, le=1.0)
    seed: "int | None" = None

    @field_validator("strength")
    @classmethod
    def _strength_bounds(cls, v: NumericOrRange) -> NumericOrRange:
        if isinstance(v, (int, float)) and not (0.0 <= float(v) <= 1.0):
            raise ValueError("strength must be within [0, 1]")
        return v


class ShadowParam(AugmentationParam):
    type: Literal["edge", "corner", "fold", "random"] = "random"


class StainParam(AugmentationParam):
    type: Literal["paper", "water", "random"] = "random"


class DiscolorationParam(AugmentationParam):
    tone: Literal["yellow", "grey", "sepia", "random"] = "random"


class CrumpleParam(AugmentationParam):
    wrinkle_density: NumericOrRange = 0.4


class JpegCompressionParam(AugmentationParam):
    min_quality: int = Field(default=35, ge=1, le=100)
    max_quality: int = Field(default=90, ge=1, le=100)


class AugmentationsConfig(BaseModel):
    """The full family of configurable document/signature augmentations."""

    model_config = ConfigDict(extra="forbid")

    rotation: AugmentationParam = Field(default_factory=lambda: AugmentationParam())
    stroke_dilation: AugmentationParam = Field(default_factory=lambda: AugmentationParam())
    motion_blur: AugmentationParam = Field(default_factory=lambda: AugmentationParam())
    gaussian_blur: AugmentationParam = Field(default_factory=lambda: AugmentationParam())
    perspective: AugmentationParam = Field(default_factory=lambda: AugmentationParam())
    scanner_noise: AugmentationParam = Field(default_factory=lambda: AugmentationParam(enabled=True, strength=0.06, probability=1.0))
    salt_pepper: AugmentationParam = Field(default_factory=lambda: AugmentationParam())
    jpeg_compression: JpegCompressionParam = Field(default_factory=lambda: JpegCompressionParam())
    contrast: AugmentationParam = Field(default_factory=lambda: AugmentationParam())
    gamma: AugmentationParam = Field(default_factory=lambda: AugmentationParam())
    coffee_stain: AugmentationParam = Field(default_factory=lambda: AugmentationParam())
    stain: StainParam = Field(default_factory=lambda: StainParam())
    shadow: ShadowParam = Field(default_factory=lambda: ShadowParam())
    crumple: CrumpleParam = Field(default_factory=lambda: CrumpleParam())
    discoloration: DiscolorationParam = Field(default_factory=lambda: DiscolorationParam())
    page_tilt: AugmentationParam = Field(default_factory=lambda: AugmentationParam())
    dust: AugmentationParam = Field(default_factory=lambda: AugmentationParam())
    scanner_streaks: AugmentationParam = Field(default_factory=lambda: AugmentationParam())
    ink_fading: AugmentationParam = Field(default_factory=lambda: AugmentationParam())
    ink_bleed: AugmentationParam = Field(default_factory=lambda: AugmentationParam())

    def items(self):
        """Yield ``(name, AugmentationParam)`` pairs for every configured effect."""
        for name in self.model_fields:
            yield name, getattr(self, name)


class PlacementConfig(BaseModel):
    """Controls where and how the signature is positioned inside its target box."""

    model_config = ConfigDict(extra="forbid")

    x_offset: NumericOrRange = 0.0
    y_offset: NumericOrRange = 0.0
    rotation_deg: NumericOrRange = 0.0
    scale_factor: NumericOrRange = 0.9
    anchor_point: Literal[
        "center", "top_left", "top_right", "bottom_left", "bottom_right",
        "top_center", "bottom_center", "left_center", "right_center",
    ] = "center"
    margin: float = Field(default=0.06, ge=0.0, le=0.45)
    misplacement_probability: float = Field(default=0.0, ge=0.0, le=1.0)
    misplacement_strength: NumericOrRange = 0.3
    allow_partial_outside: bool = False


class RenderingConfig(BaseModel):
    """Controls how the signature layer is composited onto the document."""

    model_config = ConfigDict(extra="forbid")

    blend_mode: Literal["multiply", "darken", "alpha"] = "multiply"
    ink_opacity: NumericOrRange = 0.88
    blur_sigma: NumericOrRange = 0.35
    feather_px: NumericOrRange = 1.2
    scanner_noise: NumericOrRange = 0.02
    anti_alias: bool = True


class SignerConfig(BaseModel):
    """Optional overrides for the assigned signer identity."""

    model_config = ConfigDict(extra="forbid")

    name: "str | None" = None
    role: "str | None" = None


class SignatureSelectionConfig(BaseModel):
    """Controls how a signature image is drawn from the discovered dataset."""

    model_config = ConfigDict(extra="forbid")

    source_filter: Literal["genuine", "forged", "any"] = "genuine"
    signature_id: "str | None" = None


class DocumentConfig(BaseModel):
    """Selects which template (and page) to render."""

    model_config = ConfigDict(extra="forbid")

    template: "str | None" = None
    page: "int | None" = None


class QualityConfig(BaseModel):
    """Document quality-flag inputs (feed into the ``quality`` metadata block)."""

    model_config = ConfigDict(extra="forbid")

    case_type: "str | None" = None
    notes: "str | None" = None
    scanner_quality: NumericOrRange = 0.85
    paper_quality: NumericOrRange = 0.85
    render_quality: NumericOrRange = 0.85


class DifficultyWeights(BaseModel):
    """Weights and tier thresholds for the difficulty score.

    Defaults match the specification exactly:
    ``0.35*misplacement + 0.20*whitespace + 0.20*faintness + 0.15*smallness + 0.10*low_texture``.
    """

    model_config = ConfigDict(extra="forbid")

    misplacement: float = 0.35
    whitespace: float = 0.20
    faintness: float = 0.20
    smallness: float = 0.15
    low_texture: float = 0.10
    easy_threshold: float = 0.33
    hard_threshold: float = 0.66

    @model_validator(mode="after")
    def _check_weights(self) -> "DifficultyWeights":
        total = self.misplacement + self.whitespace + self.faintness + self.smallness + self.low_texture
        if abs(total - 1.0) > 1e-6:
            raise ValueError(f"Difficulty weights must sum to 1.0, got {total}")
        return self


class PathsConfig(BaseModel):
    """Filesystem locations for inputs/outputs."""

    model_config = ConfigDict(extra="forbid")

    image_dir: str = "image"
    dataset_dir: str = "dataset"
    annotations_dir: str = "annotations"
    output_dir: str = "output"
    cache_dir: str = "output/.cache"


class GenerationConfig(BaseModel):
    """Batch-generation controls."""

    model_config = ConfigDict(extra="forbid")

    count: int = Field(default=10, ge=1)
    workers: int = Field(default=1, ge=1)
    save_visualization: bool = False
    overwrite: bool = False
    run_validation: bool = True
    start_index: int = Field(default=1, ge=0)
    filename_prefix: str = "sample"
    filename_digits: int = Field(default=6, ge=1)


class ValidationConfig(BaseModel):
    """Controls strictness of the post-generation ValidationEngine."""

    model_config = ConfigDict(extra="forbid")

    strict: bool = False
    max_coordinate_epsilon: float = 0.5


class PipelineConfig(BaseModel):
    """Top-level configuration for a full generation run."""

    model_config = ConfigDict(extra="forbid")

    seed: int = 1234
    paths: PathsConfig = Field(default_factory=PathsConfig)
    generation: GenerationConfig = Field(default_factory=GenerationConfig)
    document: DocumentConfig = Field(default_factory=DocumentConfig)
    signer: SignerConfig = Field(default_factory=SignerConfig)
    signature_selection: SignatureSelectionConfig = Field(default_factory=SignatureSelectionConfig)
    placement: PlacementConfig = Field(default_factory=PlacementConfig)
    rendering: RenderingConfig = Field(default_factory=RenderingConfig)
    augmentation: AugmentationsConfig = Field(default_factory=AugmentationsConfig)
    quality: QualityConfig = Field(default_factory=QualityConfig)
    difficulty: DifficultyWeights = Field(default_factory=DifficultyWeights)
    validation: ValidationConfig = Field(default_factory=ValidationConfig)