app/models/mapping.py

"""Pydantic models for multi-file column classification and join logic.

The AI analyzes metadata from multiple SAS files and produces:
  1. How files relate to each other (join keys)
  2. How each column should be classified (sample_id, subject_id, metadata, data, exclude)
"""

from pydantic import BaseModel, Field, field_validator


class FileJoin(BaseModel):
    """How two files should be joined."""

    left_file: str  # filename
    right_file: str  # filename
    join_columns: list[str]  # columns to join on (must exist in both files)
    join_type: str = "left"  # "left", "inner", "outer"
    reasoning: str


class ColumnClassification(BaseModel):
    """How a single source column should be used in the output."""

    source_file: str  # which file this column comes from
    source_column: str
    classification: str  # "sample_id", "subject_id", "metadata", "data", "join_key", "exclude"
    output_name: str | None = None
    confidence: float = Field(ge=0.0, le=1.0)
    reasoning: str
    transform: str | None = None


class MappingResult(BaseModel):
    """Full analysis result from the AI agent — joins + classifications."""

    session_id: str
    joins: list[FileJoin]
    columns: list[ColumnClassification]
    primary_file: str  # the main file other files join onto


class MappingRequest(BaseModel):
    """User request to analyze and classify columns across multiple files."""

    session_id: str
    user_prompt: str
    template_name: str = "aseesa_standard_v1"


class MappingConfirmation(BaseModel):
    """User-confirmed joins and column classifications."""

    session_id: str
    joins: list[FileJoin]
    columns: list[ColumnClassification]
    primary_file: str


class GenerateRequest(BaseModel):
    """Request to generate output from confirmed classification."""

    session_id: str
    output_format: str = "csv"  # "csv", "xlsx", "xpt"


# ---------------------------------------------------------------------------
# Reshape plan — for arbitrary-layout files (the raw-grid / Stars reshape flow).
#
# Instead of classifying clean columns, the agent looks at the literal cell grid
# of each sheet (see ``core/layout_extractor.py``) and proposes how to turn it
# into the Stars data matrix (variables as rows, samples as columns) plus a
# derived metadata table.
# ---------------------------------------------------------------------------


class GroupAssignment(BaseModel):
    """The cohort a sheet/block belongs to, parsed from the sheet name or a banner.

    Any field may be null if it can't be determined. These become the Sample_ID
    prefix and the derived metadata columns (Genotype/Sex/Diet).
    """

    genotype: str | None = None  # e.g. "KO", "WT", "HT"
    sex: str | None = None       # canonical "Male" / "Female"
    diet: str | None = None      # canonical "Chow" / "HFD"


class SheetReshapePlan(BaseModel):
    """How to extract one sheet/block into (sample_id, variable, value) records.

    Coordinates are 0-based indices into the rendered grid. The executor is
    tolerant: it clamps ranges and re-derives the sample columns from the header
    row, so small index errors self-heal — the semantic fields (which row holds
    sample IDs, which column holds variable labels, the group) are what matter.
    """

    logical_name: str                    # which sheet, e.g. "ITT.xlsx::KO M HFD"
    include: bool = True                 # set False to skip this sheet
    samples_axis: str = "columns"        # "columns" (samples across a row) or "rows"
    # --- columns-axis (e.g. ITT): sample IDs along a row, variables down a column ---
    sample_id_row: int | None = None     # row index holding sample IDs
    # --- rows-axis (e.g. insulin): sample labels down a column, variables along a row ---
    sample_id_col: int | None = None     # col index holding sample labels ("KO F Chow 1")
    # variable_label_index: for columns-axis this is the COLUMN of variable labels;
    #   for rows-axis it is the ROW index of the measurement-name header.
    variable_label_index: int | None = None
    # value_first/value_last: for columns-axis, first/last value ROW; for rows-axis,
    #   first/last sample ROW (both optional — the executor re-derives them).
    value_first: int | None = None
    value_last: int | None = None
    # For rows-axis side-by-side blocks (e.g. a Chow block and an HFD block on one
    # sheet): emit ONE plan entry per block, each with its own sample_id_col and its
    # own diet in `group`. The executor bounds each block's variable columns
    # automatically from sample_id_col up to the next blank header cell.
    group: GroupAssignment = GroupAssignment()
    variable_kind: str = "generic"       # "timepoint" → emit T{v} + %Basal rows; else "generic"
    notes: str = ""

    # LLMs sometimes emit null for these on excluded/edge sheets. Coerce to
    # sensible defaults so a single stray null doesn't fail the whole plan.
    @field_validator("samples_axis", mode="before")
    @classmethod
    def _default_axis(cls, v):
        return v or "columns"

    @field_validator("variable_kind", mode="before")
    @classmethod
    def _default_kind(cls, v):
        return v or "generic"

    @field_validator("group", mode="before")
    @classmethod
    def _default_group(cls, v):
        return {} if v is None else v

    @field_validator("include", mode="before")
    @classmethod
    def _default_include(cls, v):
        return True if v is None else v


class ReshapePlan(BaseModel):
    """Full agent proposal for reshaping a set of grids into Stars output."""

    sheets: list[SheetReshapePlan]
    # Canonical Sample_ID template. Tokens: {genotype} {sex} {diet} {n}
    id_format: str = "{genotype}-{sex}-{diet}-{n}"
    # Maps for shortening derived values when building the ID (full forms kept in metadata).
    sex_map: dict[str, str] = {"Male": "M", "Female": "F"}
    diet_map: dict[str, str] = {"Chow": "C", "HFD": "HFD"}
    compute_percent_basal: bool = True   # for timepoint sheets, add %Basal T{v} rows
    reasoning: str = ""


class ReshapeConfirmation(BaseModel):
    """User-confirmed reshape plan sent back from the classify page."""

    session_id: str
    plan: ReshapePlan