AI_Agro_Support/pipeline/worklist.py

"""Crop-first job discovery: find every (field, disease) pair to run today.

Unlike the original single-field pipeline, the batch does not start from a
field ID. It starts from `config.yaml`'s `crops:` list, resolves which fields
grow each crop, and for each of that crop's configured diseases finds the
fields that also have a matching disease-forecast model configured
(`AI_agrosupport_agro_models`). Each match becomes one `Job`.
"""

from __future__ import annotations

import logging
from dataclasses import dataclass

import pyodbc

from pipeline.config import CropPlan, WorklistSettings
from pipeline.db import fetch_all
from pipeline.errors import PipelineError
from pipeline.vocab import Vocabulary

logger = logging.getLogger(__name__)


@dataclass(frozen=True)
class Job:
    """One field/disease pair to run the pipeline for."""

    field_id: int
    crop_english: str
    crop_italian: str
    model_name: str  # e.g. "PERONOSPORA"; matched against anmod_model
    disease_english: str  # canonical English disease name
    organic: bool
    station: int | None
    anmod_id: int


_COLTURE_SQL = "SELECT colture_id, colture_name FROM AI_agrosupport_an_colture"

_JOBS_SQL_TEMPLATE = """
SELECT
    l.cmplay_id,
    l.cmplay_colture,
    l.cmplay_imp,
    l.cmplay_station,
    m.anmod_id,
    m.anmod_enabled
FROM AI_agrosupport_cmp_layers l
JOIN AI_agrosupport_agro_models m ON m.anmod_cmplay = l.cmplay_id
WHERE l.cmplay_colture IN ({placeholders})
  AND UPPER(LTRIM(RTRIM(m.anmod_model))) = ?
"""


def _load_crop_to_colture_ids(
    conn: pyodbc.Connection, crop_vocab: Vocabulary
) -> dict[str, list[int]]:
    """
    Map every canonical English crop name to the colture_ids that resolve to it.

    `Vocabulary.to_english` already strips whitespace and lower-cases before
    matching, so `AI_agrosupport_an_colture` rows like `"Vite "` (trailing
    space) resolve correctly without extra cleanup here. Colture names with no
    vocabulary entry are logged once and skipped, so a single unmapped crop
    cannot break discovery for every other crop.
    """
    mapping: dict[str, list[int]] = {}
    unmapped: list[str] = []
    for row in fetch_all(conn, _COLTURE_SQL):
        italian = (row.colture_name or "").strip()
        if not italian:
            continue
        try:
            english = crop_vocab.to_english(italian, kind="crop")
        except PipelineError:
            unmapped.append(italian)
            continue
        mapping.setdefault(english, []).append(int(row.colture_id))

    if unmapped:
        logger.warning(
            "%d colture name(s) have no crop vocabulary mapping and will be "
            "ignored by the worklist: %s",
            len(unmapped),
            sorted(set(unmapped)),
        )
    return mapping


def _colture_names_by_id(conn: pyodbc.Connection) -> dict[int, str]:
    return {
        int(row.colture_id): (row.colture_name or "").strip()
        for row in fetch_all(conn, _COLTURE_SQL)
    }


def build_jobs(
    conn: pyodbc.Connection,
    crops: tuple[CropPlan, ...],
    worklist: WorklistSettings,
    crop_vocab: Vocabulary,
) -> list[Job]:
    """Discover every (field, disease) job implied by `crops` and `worklist`."""
    crop_to_colture = _load_crop_to_colture_ids(conn, crop_vocab)
    colture_names = _colture_names_by_id(conn)
    allowlist = set(worklist.field_allowlist)

    jobs: list[Job] = []
    for crop_plan in crops:
        colture_ids = crop_to_colture.get(crop_plan.crop, [])
        if not colture_ids:
            logger.warning(
                "No field layer maps to crop '%s' (check vocab/crops.yaml against "
                "AI_agrosupport_an_colture); skipping its %d configured disease(s).",
                crop_plan.crop,
                len(crop_plan.diseases),
            )
            continue

        placeholders = ",".join("?" for _ in colture_ids)
        sql = _JOBS_SQL_TEMPLATE.format(placeholders=placeholders)

        for disease_plan in crop_plan.diseases:
            params = (*colture_ids, disease_plan.model_name.strip().upper())
            rows = fetch_all(conn, sql, params)

            kept = 0
            for row in rows:
                if allowlist and int(row.cmplay_id) not in allowlist:
                    continue
                if worklist.require_enabled_model and not bool(row.anmod_enabled):
                    continue

                jobs.append(
                    Job(
                        field_id=int(row.cmplay_id),
                        crop_english=crop_plan.crop,
                        crop_italian=colture_names.get(int(row.cmplay_colture), ""),
                        model_name=disease_plan.model_name,
                        disease_english=disease_plan.disease,
                        organic=row.cmplay_imp == 3,
                        station=(
                            int(row.cmplay_station)
                            if row.cmplay_station is not None
                            else None
                        ),
                        anmod_id=int(row.anmod_id),
                    )
                )
                kept += 1

            logger.info(
                "Worklist: crop=%s disease=%s (model=%s) -> %d field(s) "
                "(%d row(s) matched before allowlist/enabled filters)",
                crop_plan.crop,
                disease_plan.disease,
                disease_plan.model_name,
                kept,
                len(rows),
            )

    return jobs