scripts/estimate_macros_from_segments.py

"""Estimate meal macros from segmented area fractions and a macro lookup CSV.

This is the nutrition-side glue for the project:

1. The segmentation model predicts visible area by class.
2. Portion assumptions convert relative area to approximate grams.
3. A USDA-derived macro table converts grams to calories/macros.

Example:
    python scripts/estimate_macros_from_segments.py --segments data/nutrition/sample_segments.json
"""

from __future__ import annotations

import argparse
import csv
import json
from pathlib import Path
from typing import Any

import yaml


PROJECT_ROOT = Path(__file__).resolve().parents[1]
DEFAULT_SEGMENTS = PROJECT_ROOT / "data" / "nutrition" / "sample_segments.json"
DEFAULT_MACROS = PROJECT_ROOT / "data" / "nutrition" / "macros_per_100g.csv"
DEFAULT_PORTIONS = PROJECT_ROOT / "configs" / "portion_assumptions.yaml"
MACRO_FIELDS = ("kcal", "protein", "fat", "carbs")


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--segments", type=Path, default=DEFAULT_SEGMENTS)
    parser.add_argument("--macros", type=Path, default=DEFAULT_MACROS)
    parser.add_argument("--portions", type=Path, default=DEFAULT_PORTIONS)
    parser.add_argument("--output", type=Path)
    return parser.parse_args()


def load_segments(path: Path) -> list[dict[str, Any]]:
    payload = json.loads(path.read_text(encoding="utf-8"))
    if isinstance(payload, list):
        return payload
    if isinstance(payload, dict) and "segments" in payload:
        return payload["segments"]
    if isinstance(payload, dict):
        return [
            {"class_name": class_name, "area_fraction": area_fraction}
            for class_name, area_fraction in payload.items()
        ]
    raise ValueError(f"Unsupported segment JSON shape: {path}")


def load_macro_table(path: Path) -> dict[str, dict[str, float]]:
    with path.open(newline="", encoding="utf-8") as f:
        rows = list(csv.DictReader(f))

    table = {}
    for row in rows:
        class_name = row["class_name"]
        table[class_name] = {
            field: float(row[field]) if row.get(field) not in (None, "") else 0.0
            for field in MACRO_FIELDS
        }
    return table


def load_portions(path: Path) -> dict[str, Any]:
    return yaml.safe_load(path.read_text(encoding="utf-8"))


def estimate_grams(
    segments: list[dict[str, Any]],
    portions: dict[str, Any],
) -> dict[str, float]:
    total_plate_grams = float(portions["total_plate_grams"])
    class_config = portions.get("classes", {})

    weighted_areas = {}
    for segment in segments:
        class_name = str(segment["class_name"])
        area_fraction = float(segment["area_fraction"])
        density_weight = float(class_config.get(class_name, {}).get("density_weight", 1.0))
        weighted_areas[class_name] = weighted_areas.get(class_name, 0.0) + area_fraction * density_weight

    total_weighted_area = sum(weighted_areas.values())
    if total_weighted_area <= 0:
        raise ValueError("Segment area fractions must sum to a positive value.")

    return {
        class_name: total_plate_grams * weighted_area / total_weighted_area
        for class_name, weighted_area in weighted_areas.items()
    }


def estimate_macros(
    grams_by_class: dict[str, float],
    macro_table: dict[str, dict[str, float]],
) -> tuple[list[dict[str, Any]], dict[str, float]]:
    rows = []
    totals = {field: 0.0 for field in ("grams", *MACRO_FIELDS)}

    for class_name, grams in grams_by_class.items():
        if class_name not in macro_table:
            raise KeyError(f"Missing macro row for class '{class_name}'")

        row = {"class_name": class_name, "grams": round(grams, 1)}
        for field in MACRO_FIELDS:
            value = grams / 100.0 * macro_table[class_name][field]
            row[field] = round(value, 1)
            totals[field] += value

        totals["grams"] += grams
        rows.append(row)

    totals = {key: round(value, 1) for key, value in totals.items()}
    return rows, totals


def main() -> None:
    args = parse_args()
    segments = load_segments(args.segments)
    macro_table = load_macro_table(args.macros)
    portions = load_portions(args.portions)

    grams_by_class = estimate_grams(segments, portions)
    rows, totals = estimate_macros(grams_by_class, macro_table)
    payload = {"items": rows, "totals": totals}

    if args.output:
        args.output.parent.mkdir(parents=True, exist_ok=True)
        args.output.write_text(json.dumps(payload, indent=2), encoding="utf-8")
        print(f"Wrote estimate: {args.output}")
    else:
        print(json.dumps(payload, indent=2))


if __name__ == "__main__":
    main()