chain_problems.py

# chain_problems.py
import json
import logging
from typing import Dict, Optional
from langchain import PromptTemplate, LLMChain
from models import chat_model

logger = logging.getLogger(__name__)

problem_prompt_template = PromptTemplate(
    input_variables=["responses", "internal_report"],
    template=(
        "You are a wellness analyst. You have the following user responses to health-related questions:\n"
        "{responses}\n\n"
        "You also have an internal analysis report:\n"
        "{internal_report}\n\n"
        "From these inputs, determine a 'problem severity percentage' for the user in the following areas: "
        "stress_management, low_therapy, balanced_weight, restless_night, lack_of_motivation, gut_health, anxiety, burnout.\n\n"
        "Consider how these severity percentages will later be used to recommend suitable wellness packages such as:\n\n"
        "1. Fitness Mobility\n"
        "2. Balance Weight\n"
        "3. No More Insomnia\n"
        "4. Chronic Care\n"
        "5. Mental Wellness\n"
        "6. Focus Flow\n\n"
        "Return your answer in JSON format with keys: stress_management, low_therapy, balanced_weight, restless_night, "
        "lack_of_motivation, gut_health, anxiety, burnout.\n"
        "Ensure severity percentages are numbers from 0 to 100.\n\n"
        "JSON Output:"
    )
)

problem_chain = LLMChain(llm=chat_model, prompt=problem_prompt_template)


def recommend_wellness_packages(
    problems: Dict[str, float],
    sleep_hours: Optional[float] = None
) -> Dict[str, str]:
    """
    Given the severity dictionary and the user's reported sleep hours,
    recommend one or more wellness packages. 

    Returns:
        A dictionary with:
          - 'recommended_packages': comma-separated string of package names
          - 'recommendation_details': textual explanation of why these were chosen
    """
    # Threshold for considering an issue "significant"
    THRESHOLD = 50.0
    
    recommended = []
    rationale = []

    # 1. Stress / Mood => "Mental Wellness", "Focus Flow"
    if (
        problems["stress_management"] >= THRESHOLD 
        or problems["anxiety"] >= THRESHOLD 
        or problems["burnout"] >= THRESHOLD
        or problems["low_therapy"] >= THRESHOLD
        or problems["lack_of_motivation"] >= THRESHOLD
    ):
        recommended.append("Mental Wellness")
        rationale.append(
            "Mental Wellness Package suggested due to elevated stress, anxiety, "
            "or motivation-related concerns."
        )
    
    if problems["lack_of_motivation"] >= THRESHOLD or problems["stress_management"] >= THRESHOLD:
        recommended.append("Focus Flow")
        rationale.append(
            "Focus Flow Package recommended to improve concentration and reduce mental fatigue."
        )

    # 2. Balanced Weight / Gut Health => "Balance Weight"
    if problems["balanced_weight"] >= THRESHOLD or problems["gut_health"] >= THRESHOLD:
        recommended.append("Balance Weight")
        rationale.append(
            "Balance Weight Package suggested due to weight management or gut health concerns."
        )

    # 3. Restless nights => "No More Insomnia"
    #    **FIX**: Only recommend if user sleeps < 4 hours (to avoid recommending it for 10+ hours).
    if sleep_hours is not None and sleep_hours < 4:
        # If the user is sleeping less than 4 hours, we consider insomnia a big issue.
        recommended.append("No More Insomnia")
        rationale.append(
            "No More Insomnia Package recommended because the user reported sleeping less than 4 hours."
        )

    # 4. High burnout => "Chronic Care"
    if problems["burnout"] >= THRESHOLD:
        recommended.append("Chronic Care")
        rationale.append(
            "Chronic Care Package recommended to help manage long-term stress and burnout."
        )

    # 5. If user has borderline or higher weight/stress => "Fitness Mobility"
    if (
        problems["balanced_weight"] >= (THRESHOLD * 0.7) 
        or problems["stress_management"] >= (THRESHOLD * 0.7)
    ):
        recommended.append("Fitness Mobility")
        rationale.append(
            "Fitness Mobility Package may support better physical conditioning and relieve stress."
        )

    # Ensure uniqueness if appended multiple times
    unique_packages = list(dict.fromkeys(recommended))

    # Build a clear explanation
    recommendation_details = "\n".join(rationale) if rationale else "No particular issues detected."

    return {
        "recommended_packages": ", ".join(unique_packages),
        "recommendation_details": recommendation_details
    }


def analyze_problems_with_chain(responses: Dict[str, str], internal_report: str) -> Dict[str, object]:
    """
    Analyzes user responses and internal report to extract problem severity
    and then recommends relevant wellness packages.
    
    Returns:
        A dictionary with:
          - keys for each problem severity (stress_management, low_therapy, etc.)
          - 'recommended_packages': comma-separated package names
          - 'recommendation_details': textual explanation
    """
    responses_str = "\n".join(f"{q}: {a}" for q, a in responses.items())
    raw_text = problem_chain.run(responses=responses_str, internal_report=internal_report)
    
    # Try to parse problem severities from LLM output
    try:
        start_idx = raw_text.find('{')
        end_idx = raw_text.rfind('}') + 1
        json_str = raw_text[start_idx:end_idx]
        problems = json.loads(json_str)

        # Ensure all eight keys exist
        for key in [
            "stress_management",
            "low_therapy",
            "balanced_weight",
            "restless_night",
            "lack_of_motivation",
            "gut_health",
            "anxiety",
            "burnout"
        ]:
            problems.setdefault(key, 0.0)

        # Convert values to float
        problems = {k: float(v) for k, v in problems.items()}

    except Exception as e:
        logger.error(f"Error parsing problem percentages from LLM: {e}")
        # Default to zero severities
        problems = {
            "stress_management": 0.0,
            "low_therapy": 0.0,
            "balanced_weight": 0.0,
            "restless_night": 0.0,
            "lack_of_motivation": 0.0,
            "gut_health": 0.0,
            "anxiety": 0.0,
            "burnout": 0.0
        }

    # --- EXTRACT USER SLEEP HOURS ---
    # We assume the user answer is in responses["sleep_duration"] or similar key.
    # If your data uses a different key for the hours of sleep, adjust accordingly.
    sleep_hours = None
    if "sleep_duration" in responses:
        try:
            sleep_hours = float(responses["sleep_duration"])
        except ValueError:
            sleep_hours = None

    # Now get the recommendations based on problem severities + sleep_hours
    recommendations = recommend_wellness_packages(problems, sleep_hours=sleep_hours)

    # Merge and return
    return {
        **problems,
        **recommendations
    }