engine.py

from __future__ import annotations

from typing import Final
from schema import (
    ActivityLevel, LifeStage, NeuteredSpayed,
    NutritionTargets, PetProfile, PetType, ReproStatus,
)


# MER multipliers 
_ENERGY_MULTIPLIERS: Final[dict[str, dict[str, float]]] = {
    "dog": {
        "adult_intact":        1.8,
        "adult_neutered":      1.6,
        "adult_obesity_prone": 1.4,
        "puppy_over_4m":       2.0,
        "pregnant":            1.8,
        "lactating_peak":      3.0,
        "senior_active":       1.6,
        "senior_sedentary":    1.4,
    },
    "cat": {
        "adult_intact":        1.4,
        "adult_neutered":      1.2,
        "adult_obesity_prone": 1.0,
        "kitten":              2.5,
        "pregnant":            1.4,
        "lactating_peak":      2.5,
        "senior_active":       1.2,
        "senior_sedentary":    1.1,
    },
}

# Macronutrient targets — % of daily calories 
_MACRONUTRIENT_TARGETS: Final[dict[str, dict[str, dict[str, float]]]] = {
    "dog": {
        "adult_maintenance": {"protein_pct": 20.0, "fat_pct": 25.0, "min_protein_pct": 18.0},
        "puppy_growth":      {"protein_pct": 25.0, "fat_pct": 30.0, "min_protein_pct": 22.5},
        "senior":            {"protein_pct": 22.0, "fat_pct": 20.0, "min_protein_pct": 20.0},
    },
    "cat": {
        "adult_maintenance": {"protein_pct": 30.0, "fat_pct": 30.0, "min_protein_pct": 26.0},
        "kitten_growth":     {"protein_pct": 35.0, "fat_pct": 35.0, "min_protein_pct": 30.0},
        "senior":            {"protein_pct": 32.0, "fat_pct": 25.0, "min_protein_pct": 28.0},
    },
}

# Atwater energy densities (kcal/g) 
_ENERGY_DENSITIES: Final[dict[str, float]] = {
    "protein":      4.0,
    "carbohydrate": 4.0,
    "fat":          9.0,
}

_DIGESTIBILITY_FACTOR: Final[float] = 0.90  
_WATER_ML_PER_KCAL:    Final[float] = 1.0   

_WATER_MULTIPLIERS: Final[dict[str, float]] = {
    "lactating":     2.0,
    "pregnant":      1.5,
    "puppy":         1.5,
    "kitten":        1.5,
    "senior":        1.2,
    "high_activity": 1.5,
}

_WEIGHT_RANGES: Final[dict[str, tuple[float, float]]] = {
    "dog": (0.5, 90.0),
    "cat": (2.0, 12.0),
}

class NutritionEngine:

    def calculate(self, profile: PetProfile) -> NutritionTargets:
        lo, hi = _WEIGHT_RANGES[profile.pet_type]
        if not lo <= profile.weight_kg <= hi:
            raise ValueError(
                f"Weight {profile.weight_kg} kg is outside the plausible range "
                f"({lo}–{hi} kg) for a {profile.pet_type}."
            )
        rer        = 70.0 * (profile.weight_kg ** 0.75)
        multiplier = self._mer_multiplier(profile)
        mer        = round(rer * multiplier, 2)
        macros     = self._macros(mer, profile)
        water      = self._water(mer, profile)

        return NutritionTargets(
            daily_calorie_target  = mer,
            daily_protein_g       = round(macros["protein_g"],       2),
            daily_fat_g           = round(macros["fat_g"],           2),
            daily_carbohydrates_g = round(macros["carbohydrates_g"], 2),
            daily_water_ml        = round(water,                     2),
        )

    @staticmethod
    def _mer_multiplier(p: PetProfile) -> float:
        table = _ENERGY_MULTIPLIERS[p.pet_type]
        if p.pregnant_lactating == ReproStatus.LACTATING:
            key = "lactating_peak"
        elif p.pregnant_lactating == ReproStatus.PREGNANT:
            key = "pregnant"
        elif p.life_stage == LifeStage.PUPPY:
            key = "puppy_over_4m"
        elif p.life_stage == LifeStage.KITTEN:
            key = "kitten"
        elif p.life_stage == LifeStage.SENIOR:
            key = "senior_sedentary" if p.activity_level == ActivityLevel.LOW else "senior_active"
        elif p.bcs > 6:
            key = "adult_obesity_prone"
        elif p.neutered_spayed == NeuteredSpayed.YES:
            key = "adult_neutered"
        else:
            key = "adult_intact"
        return table.get(key, 1.6)

    @staticmethod
    def _macro_key(p: PetProfile) -> str:
        if p.life_stage == LifeStage.SENIOR:
            return "senior"
        if p.life_stage == LifeStage.PUPPY:
            return "puppy_growth"
        if p.life_stage == LifeStage.KITTEN:
            return "kitten_growth"
        return "adult_maintenance"

    @classmethod
    def _macros(cls, kcal: float, p: PetProfile) -> dict[str, float]:
        tbl = (
            _MACRONUTRIENT_TARGETS[p.pet_type].get(cls._macro_key(p))
            or _MACRONUTRIENT_TARGETS[p.pet_type]["adult_maintenance"]
        )
        prot_pct = tbl["protein_pct"]
        fat_pct  = tbl["fat_pct"]

        if p.pregnant_lactating == ReproStatus.PREGNANT:
            prot_pct *= 1.25
        elif p.pregnant_lactating == ReproStatus.LACTATING:
            prot_pct *= 1.50

        prot_pct = max(prot_pct, tbl["min_protein_pct"])

        p_kcal = kcal * prot_pct / 100
        f_kcal = kcal * fat_pct  / 100
        c_kcal = kcal - p_kcal - f_kcal
        d      = _DIGESTIBILITY_FACTOR

        return {
            "protein_g":       p_kcal / _ENERGY_DENSITIES["protein"]      * d,
            "fat_g":           f_kcal / _ENERGY_DENSITIES["fat"]           * d,
            "carbohydrates_g": c_kcal / _ENERGY_DENSITIES["carbohydrate"],
        }

    @staticmethod
    def _water(kcal: float, p: PetProfile) -> float:
        base = kcal * _WATER_ML_PER_KCAL
        multiplier = (
            _WATER_MULTIPLIERS["lactating"]     if p.pregnant_lactating == ReproStatus.LACTATING else
            _WATER_MULTIPLIERS["pregnant"]      if p.pregnant_lactating == ReproStatus.PREGNANT  else
            _WATER_MULTIPLIERS["puppy"]         if p.life_stage == LifeStage.PUPPY               else
            _WATER_MULTIPLIERS["kitten"]        if p.life_stage == LifeStage.KITTEN              else
            _WATER_MULTIPLIERS["senior"]        if p.life_stage == LifeStage.SENIOR              else
            _WATER_MULTIPLIERS["high_activity"] if p.activity_level == ActivityLevel.HIGH        else
            1.0
        )
        return base * multiplier