from dspy.teleprompt import BootstrapFewShotWithRandomSearch

from app.ai.nutrition import nutrition_module
from app.core.ai_config import configure_dspy
from scripts.nutrition_data import train_examples

# 0. Configure DSPy
configure_dspy()


# 1. Define Advanced Metric
def validate_nutrition_v2(example, pred, trace=None):
    # Condition A: Accuracy (within 15% of ground truth)
    actual_cals = example.nutritional_info.calories
    pred_cals = pred.nutritional_info.calories

    threshold = 0.15
    lower = actual_cals * (1 - threshold)
    upper = actual_cals * (1 + threshold)
    is_accurate_count = lower <= pred_cals <= upper

    # Condition B: Consistency (Macros match Calories within 20%)
    # This prevents "hallucinated" numbers that don't satisfy physics
    p = pred.nutritional_info.protein
    c = pred.nutritional_info.carbs
    f = pred.nutritional_info.fats

    calculated_cals = (p * 4) + (c * 4) + (f * 9)
    # Using a slightly looser bounds (20%) for fiber/rounding
    consistency_threshold = 0.20
    is_consistent_math = abs(calculated_cals - pred_cals) < (pred_cals * consistency_threshold)

    # We want BOTH to be true
    return is_accurate_count and is_consistent_math


# 2. Setup Advanced Optimizer
# RandomSearch is more expensive but finds better reasoning traces by randomizing
# the selection of few-shot examples.
# num_candidate_programs=10 means it will try 10 different combinations of prompts/examples
print("Configuring RandomSearch Optimizer...")
teleprompter = BootstrapFewShotWithRandomSearch(
    metric=validate_nutrition_v2,
    max_bootstrapped_demos=4,
    max_labeled_demos=4,
    num_candidate_programs=5,  # Reduced to 5 for speed in this demo, typically 10-20
    num_threads=1,  # Sequential for stability, increase for parallelism
)

# 3. Compile (Optimize) the Module
print("Optimizing V2 (this includes random search and macro checks)...")
# Note: assertions are compiled into the pipeline automatically in newer DSPy,
# acting as soft constraints during the search.
compiled_nutrition = teleprompter.compile(nutrition_module, trainset=train_examples)

# 4. Save
compiled_nutrition.save("app/ai/nutrition_compiled.json")
print("Optimization V2 complete! Overwrote app/ai/nutrition_compiled.json")