services/models.py

"""
Pydantic schemas shared by the FastAPI endpoints.
"""

from typing import List, Optional

from pydantic import BaseModel, Field


class BasePlotRequest(BaseModel):
    """Fields common to every plot request.

    hf_token is deliberately NOT part of this schema: it's a write-capable
    credential, so it's read server-side from the HF_TOKEN Space secret
    (see app.py's _server_hf_token()) instead of being accepted from the
    caller. This lets a public webpage call these endpoints directly
    without ever handling the token.
    """

    dataset_repo: str = Field(..., description="HF dataset repo, e.g. 'lopezels/public-pdfs'")
    grid_name: str = Field(..., description="LHAPDF set name, e.g. 'CT25NNLO'")
    # If False, always recompute even if a matching cached plot exists.
    use_cache: bool = True


class StandardPlotRequest(BasePlotRequest):
    """Central value + uncertainty band for x*f(x, Q) of a single flavor."""

    q_scale: float = 91.2
    parton_id: int = 21
    plot_color: Optional[str] = Field(None, description="Hex color; defaults to the flavor's palette color")
    custom_title: Optional[str] = None
    x_min: float = 1e-5
    x_max: float = 10 ** -0.01
    n_points: int = 150


class RatioPlotRequest(BasePlotRequest):
    """Same as StandardPlotRequest, normalized to the central value (ratio to central = 1)."""

    q_scale: float = 91.2
    parton_id: int = 21
    plot_color: Optional[str] = None
    custom_title: Optional[str] = None
    x_min: float = 1e-5
    x_max: float = 10 ** -0.01
    n_points: int = 150


class CorrelationPlotRequest(BasePlotRequest):
    """Hessian-style correlation contour between two flavors."""

    q_scale: float = 91.2
    parton_id_1: int = 2
    parton_id_2: int = 21
    color_correlated: Optional[str] = Field(None, description="Defaults to parton_id_1's palette color")
    color_anti_correlated: str = "black"
    x_min: float = 1e-4
    x_max: float = 10 ** -0.01
    n_points: int = 35


class MainPlotRequest(BasePlotRequest):
    """
    The classic multi-flavor "overview" plot (CT18 Fig. 2 style): every
    flavor in parton_ids overlaid on one figure at a single Q. The gluon
    (id 21) is always scaled down by a factor of 5 in the figure itself
    (see services.plotting.build_main_plot), not configurable here.
    """

    q_scale: float = 2.0
    parton_ids: List[int] = [3, 21, 2, 1, -1, -2, 4]
    custom_title: Optional[str] = None
    x_min: float = 1e-6
    x_max: float = 0.9
    n_points: int = 500


class PlotResponse(BaseModel):
    status: str
    cached: bool
    message: str
    plot_id: str
    plot_url: str
    plotly_json_url: str