app.py

import os
from pathlib import Path
from typing import Optional, Tuple, List, Dict

import gradio as gr
import pandas as pd
import numpy as np
import plotly.express as px
import joblib

# ZeroGPU hooks (safe on CPU Spaces too)
import spaces
import torch

# Optional micro-model to "polish" text when GPU window is available
from transformers import AutoModelForSeq2SeqLM, AutoTokenizer, pipeline

# ---------------------
# Constants & storage
# ---------------------
DATA_DIR = Path("data"); DATA_DIR.mkdir(exist_ok=True)
TS_FMT = "%Y-%m-%d %H:%M:%S"

# Load your regressor
DT_PATH = "./decision_tree_regressor.joblib"
decision_tree_regressor = joblib.load(DT_PATH)

# Lightweight text model (CPU ok, faster on GPU)
GEN_MODEL = os.getenv("PLAN_POLISH_MODEL", "google/flan-t5-small")
_tokenizer = AutoTokenizer.from_pretrained(GEN_MODEL)
_model = AutoModelForSeq2SeqLM.from_pretrained(GEN_MODEL)
_generate_cpu = pipeline("text2text-generation", model=_model, tokenizer=_tokenizer, device=-1)

# --------------
# ZeroGPU fns
# --------------
@spaces.GPU
def gpu_warmup() -> str:
    return f"cuda={torch.cuda.is_available()}"

@spaces.GPU
def polish_on_gpu(text: str, lang: str = "en") -> str:
    """Polish/translate the already-generated plan inside a GPU window.
    Falls back to CPU gracefully if needed.
    """
    try:
        if torch.cuda.is_available():
            gen = pipeline(
                "text2text-generation",
                model=_model.to("cuda"),
                tokenizer=_tokenizer,
                device=0,
            )
        else:
            gen = _generate_cpu
        prompt = (
            "Rewrite the following fasting plan in a friendly coaching tone, keep markdown structure, "
            f"and output language '{lang}'. Keep tables and numbered lists concise.\n\n" + text
        )
        out = gen(prompt, max_new_tokens=700)
        return out[0]["generated_text"].strip()
    except Exception as e:
        out = _generate_cpu(text, max_new_tokens=10)
        return text + f"\n\n(Polish step skipped: {e})"

try:
    _ = gpu_warmup()
except Exception:
    pass

# ---------------------
# Utilities (metrics)
# ---------------------
ACTIVITY = {
    "Sedentary": 1.2,
    "Lightly active": 1.375,
    "Moderately active": 1.55,
    "Very active": 1.725,
    "Athlete": 1.9,
}

GOAL_CAL_ADJ = {  # % change to TDEE
    "Fat loss": -0.15,
    "Recomp/Maintenance": 0.0,
    "Muscle gain": 0.10,
}

def bmi(weight_kg: float, height_cm: float) -> float:
    return weight_kg / ((height_cm / 100) ** 2)


def bmr_mifflin(sex: str, weight_kg: float, height_cm: float, age: float) -> float:
    s = 5 if sex == "Male" else -161
    return 10 * weight_kg + 6.25 * height_cm - 5 * age + s


def tdee(bmr: float, activity: str) -> float:
    return bmr * ACTIVITY.get(activity, 1.2)


def parse_hhmm(hhmm: str) -> Tuple[int, int]:
    h, m = hhmm.split(":")
    h = int(h); m = int(m)
    if not (0 <= h <= 23 and 0 <= m <= 59):
        raise ValueError("Time must be HH:MM in 24h format.")
    return h, m


def fmt_hhmm(h: int, m: int) -> str:
    return f"{h:02d}:{m:02d}"

# ---------------------
# Plan generator (deterministic, rich)
# ---------------------
DIET_STYLES = ["Omnivore", "Mediterranean", "Vegetarian", "Vegan", "Low-carb"]

MEAL_IDEAS = {
    "Omnivore": [
        "Greek yogurt + berries + nuts",
        "Chicken bowl (rice, veggies, olive oil)",
        "Eggs, avocado, sourdough",
        "Salmon, quinoa, asparagus",
        "Lean beef, sweet potato, salad",
        "Tuna whole-grain wrap",
        "Cottage cheese + fruit + seeds",
    ],
    "Mediterranean": [
        "Oats with dates, walnuts, olive oil drizzle",
        "Grilled fish, lentil salad, greens",
        "Hummus platter, wholegrain pita, veg",
        "Chickpea tomato stew",
        "Feta + olive salad, quinoa",
        "Shakshuka + side salad",
        "Lentils, roasted veg, tahini",
    ],
    "Vegetarian": [
        "Tofu scramble, toast, avocado",
        "Paneer tikka bowl",
        "Bean chili + brown rice",
        "Halloumi, couscous, veg",
        "Greek salad + eggs",
        "Tempeh stir-fry",
        "Yogurt parfait + granola",
    ],
    "Vegan": [
        "Tofu scramble, avocado toast",
        "Lentil curry + basmati",
        "Burrito bowl (beans, corn, salsa)",
        "Seitan, roasted potatoes, veg",
        "Tofu poke bowl",
        "Chickpea pasta + marinara",
        "Overnight oats + banana + peanut butter",
    ],
    "Low-carb": [
        "Eggs, smoked salmon, salad",
        "Chicken Caesar (no croutons)",
        "Beef & greens stir-fry",
        "Omelette + veg + cheese",
        "Zoodles + turkey bolognese",
        "Tofu salad w/ tahini",
        "Yogurt + nuts (moderate)",
    ],
}

WORKOUTS = {
    "Fat loss": [
        "3× LISS cardio 30–40min",
        "2× full‑body strength 45min",
        "1× intervals 12–16min",
        "Daily 8–10k steps"
    ],
    "Recomp/Maintenance": [
        "3× full‑body strength 45–60min",
        "1–2× LISS cardio 30min",
        "Mobility 10min daily",
        "8–10k steps"
    ],
    "Muscle gain": [
        "4× strength split 45–60min",
        "Optional 1× LISS 20–30min",
        "Mobility 10min",
        "7–9k steps"
    ],
}


def feeding_schedule(first_meal_hhmm: str, fasting_hours: float) -> List[Tuple[str, str]]:
    """Return 7 (start,end) strings for the eating window each day."""
    h, m = parse_hhmm(first_meal_hhmm)
    window = max(0.0, 24 - float(fasting_hours))
    start_minutes = h * 60 + m
    end_minutes = int((start_minutes + window * 60) % (24 * 60))

    sched = []
    for _ in range(7):
        start = fmt_hhmm(h, m)
        end = fmt_hhmm(end_minutes // 60, end_minutes % 60)
        sched.append((start, end))
    return sched


def weekly_plan(diet: str, sched: List[Tuple[str, str]], kcal: int, protein_g: int) -> pd.DataFrame:
    ideas = MEAL_IDEAS[diet]
    rows = []
    for i in range(7):
        day = ["Mon","Tue","Wed","Thu","Fri","Sat","Sun"][i]
        start, end = sched[i]
        meal1 = ideas[i % len(ideas)]
        meal2 = ideas[(i+3) % len(ideas)]
        snack = "Fruit or nuts (optional)"
        rows.append({
            "Day": day,
            "Feeding window": f"{start}–{end}",
            "Meal 1": meal1,
            "Meal 2": meal2,
            "Protein target": f"≥ {protein_g} g",
            "Daily kcal": kcal,
            "Snack": snack,
        })
    return pd.DataFrame(rows)


def shopping_list(diet: str) -> List[str]:
    core = [
        "Leafy greens, mixed veg, berries",
        "Olive oil, nuts/seeds, herbs & spices",
        "Coffee/tea, mineral water, electrolytes",
    ]
    extras = {
        "Omnivore": ["Chicken, fish, eggs, yogurt, cottage cheese", "Rice/quinoa/sourdough", "Beans/lentils"],
        "Mediterranean": ["Fish, feta, olives", "Whole grains (bulgur, farro)", "Chickpeas/lentils"],
        "Vegetarian": ["Eggs, dairy, paneer", "Legumes", "Tofu/tempeh"],
        "Vegan": ["Tofu/tempeh/seitan", "Beans/lentils", "Plant yogurt/milk"],
        "Low-carb": ["Eggs, fish, meat", "Green veg", "Greek yogurt, cheese"],
    }
    return core + extras[diet]

# ---------------------
# Tracker (history)
# ---------------------
active_fasts: Dict[str, pd.Timestamp] = {}

def _csv(u: str) -> Path:
    safe = "".join(ch for ch in (u or "default") if ch.isalnum() or ch in ("_","-"))
    return DATA_DIR / f"{safe}.csv"

def hist_load(u: str) -> pd.DataFrame:
    p = _csv(u)
    if p.exists():
        d = pd.read_csv(p)
        for c in ["start_time","end_time"]:
            if c in d: d[c] = pd.to_datetime(d[c], errors="coerce")
        return d
    return pd.DataFrame(columns=["start_time","end_time","duration_hours","note"])

def hist_save(u: str, d: pd.DataFrame):
    d.to_csv(_csv(u), index=False)

# ---------------------
# Core actions
# ---------------------

def predict_and_plan(fasting_duration, meal_timing, weight, age, gender, height,
                     activity, goal, diet, lang, ai_polish) -> Tuple[Optional[float], str, str, pd.DataFrame, object, str]:
    try:
        # Validation
        if fasting_duration < 0 or fasting_duration > 72:
            raise ValueError("Fasting duration must be 0–72h.")
        parse_hhmm(meal_timing)
        if weight <= 0 or height <= 0 or age < 0:
            raise ValueError("Check weight/height/age values.")

        # Model score
        df = pd.DataFrame({
            "Fasting Duration (hours)": [float(fasting_duration)],
            "Meal Timing (hour:minute)": [lambda t=meal_timing: int(t.split(":")[0]) + int(t.split(":")[1]) / 60.0][0](),
            "Body Weight (kg)": [float(weight)],
            "Age (years)": [float(age)],
            "Height (cm)": [float(height)],
            "Gender_Male": [1 if gender == "Male" else 0],
            "Gender_Other": [1 if gender == "Other" else 0],
        })
        score = float(decision_tree_regressor.predict(df)[0])

        # Metrics
        bmr = bmr_mifflin(gender, weight, height, age)
        tdee_kcal = tdee(bmr, activity)
        adj = GOAL_CAL_ADJ[goal]
        target_kcal = int(round(tdee_kcal * (1 + adj)))
        protein_g = int(round(max(1.6 * weight, 90 if goal == "Muscle gain" else 80)))
        bmi_val = round(bmi(weight, height), 1)

        # Schedule & tables
        sched = feeding_schedule(meal_timing, float(fasting_duration))
        plan_df = weekly_plan(diet, sched, target_kcal, protein_g)

        # Chart (Gantt-style feeding window)
        chart_df = pd.DataFrame({
            "Day": ["Mon","Tue","Wed","Thu","Fri","Sat","Sun"],
            "start": [int(s.split(":")[0])*60 + int(s.split(":")[1]) for s,_ in sched],
            "length": [max(0, int((24 - float(fasting_duration))*60))]*7,
        })
        fig = px.bar(chart_df, y="Day", x="length", base="start", orientation="h", title="Feeding window each day (minutes)")
        fig.update_layout(xaxis=dict(range=[0,1440], tickvals=[0,360,720,1080,1440], ticktext=["00:00","06:00","12:00","18:00","24:00"]))

        # Markdown plan
        hdr = {
            "en": "## Your 7‑day intermittent fasting plan",
            "es": "## Tu plan de ayuno intermitente de 7 días",
        }[lang]
        kpis = (
            f"**Score:** {score:.1f}  •  **BMI:** {bmi_val}  •  **BMR:** {int(bmr)} kcal  •  **TDEE:** {int(tdee_kcal)} kcal  •  "
            f"**Target:** {target_kcal} kcal  •  **Protein:** ≥ {protein_g} g  •  **Diet:** {diet}\n"
        )
        sched_md = "\n".join([f"- **{d}**: {s} – {e}" for d,(s,e) in zip(["Mon","Tue","Wed","Thu","Fri","Sat","Sun"], sched)])
        workouts = "\n".join([f"- {w}" for w in WORKOUTS[goal]])
        shop = "\n".join([f"- {x}" for x in shopping_list(diet)])

        plan_md = f"""
{hdr}

{kpis}

### Feeding window (daily)
{sched_md}

### Weekly training
{workouts}

### Daily meals (example week)
(See table below for details.)

### Shopping list
{shop}

> Hydration & electrolytes during the fast, protein at each meal, whole foods, and 7–9 hours sleep.
""".strip()

        # Optional AI polish (ZeroGPU window)
        if ai_polish:
            try:
                plan_md = polish_on_gpu(plan_md, lang)
            except Exception:
                pass

        # Export file path (Markdown)
        md_path = DATA_DIR / "plan.md"
        md_path.write_text(plan_md, encoding="utf-8")

        return score, kpis, plan_md, plan_df, fig, str(md_path)
    except Exception as e:
        return None, "", f"⚠️ {e}", pd.DataFrame(), None, ""

# ---------------------
# Tracker actions
# ---------------------

def start_fast(user: str, note: str):
    if not user: return "Enter username in Settings.", None
    if user in active_fasts: return f"Already fasting since {active_fasts[user]}.", None
    active_fasts[user] = pd.Timestamp.now()
    return f"✅ Fast started at {active_fasts[user].strftime(TS_FMT)}.", None


def end_fast(user: str):
    if not user: return "Enter username in Settings.", None, None, None
    if user not in active_fasts: return "No active fast.", None, None, None
    end = pd.Timestamp.now(); start = active_fasts.pop(user)
    dur = round((end - start).total_seconds()/3600, 2)
    df = hist_load(user)
    df.loc[len(df)] = [start, end, dur, ""]
    hist_save(user, df)
    chart = make_hist_chart(df)
    return f"✅ Fast ended at {end.strftime(TS_FMT)} • {dur} h", df.tail(12), chart, hist_stats(df)


def refresh_hist(user: str):
    df = hist_load(user)
    return df.tail(12), make_hist_chart(df), hist_stats(df)


def make_hist_chart(df: pd.DataFrame):
    if df.empty: return None
    d = df.dropna(subset=["end_time"]).copy()
    d["date"] = pd.to_datetime(d["end_time"]).dt.date
    fig = px.bar(d, x="date", y="duration_hours", title="Fasting duration by day (h)")
    fig.update_layout(height=300, margin=dict(l=10,r=10,t=40,b=10))
    return fig


def hist_stats(df: pd.DataFrame) -> str:
    if df.empty: return "No history yet."
    last7 = df.tail(7)
    avg = last7["duration_hours"].mean()
    streak = compute_streak(df)
    return f"Total fasts: {len(df)}\nAvg (last 7): {avg:.2f} h\nCurrent streak: {streak} day(s)"


def compute_streak(df: pd.DataFrame) -> int:
    d = df.dropna(subset=["end_time"]).copy()
    if d.empty: return 0
    days = set(pd.to_datetime(d["end_time"]).dt.date)
    cur = pd.Timestamp.now().date(); streak=0
    while cur in days:
        streak+=1; cur = cur - pd.Timedelta(days=1)
    return streak

# ---------------------
# UI
# ---------------------
with gr.Blocks(
    title="Intermittent Fasting Coach — Pro",
    theme=gr.themes.Soft(primary_hue=gr.themes.colors.orange, neutral_hue=gr.themes.colors.gray),
) as demo:
    gr.Markdown("""
    # 🥣 Intermittent Fasting — Pro
    Detailed coaching plans + tracker. ZeroGPU‑ready (with CPU fallback). All data stored locally in this Space.
    """)

    with gr.Tabs():
        # --- Coach tab
        with gr.TabItem("Coach"):
            with gr.Row():
                with gr.Column():
                    fasting_duration = gr.Number(label="Fasting Duration (h)", value=16, minimum=0, maximum=72, step=0.5)
                    meal_timing = gr.Textbox(label="First meal time (HH:MM)", value="12:30")
                    weight = gr.Number(label="Body Weight (kg)", value=70, step=0.5)
                with gr.Column():
                    age = gr.Slider(label="Age (years)", minimum=18, maximum=100, value=35)
                    gender = gr.Radio(["Male","Female","Other"], label="Gender", value="Male")
                    height = gr.Number(label="Height (cm)", value=175)
            with gr.Row():
                activity = gr.Dropdown(choices=list(ACTIVITY.keys()), value="Lightly active", label="Activity")
                goal = gr.Dropdown(choices=list(GOAL_CAL_ADJ.keys()), value="Recomp/Maintenance", label="Goal")
                diet = gr.Dropdown(choices=DIET_STYLES, value="Mediterranean", label="Diet style")
                lang = gr.Radio(["en","es"], value="en", label="Language")
                ai_polish = gr.Checkbox(value=True, label="AI polish (uses ZeroGPU)")

            btn = gr.Button("Predict & Build Plan", variant="primary")

            score_out = gr.Number(label="Predicted score")
            kpi_out = gr.Markdown()
            plan_md = gr.Markdown()
            plan_tbl = gr.Dataframe(headers=["Day","Feeding window","Meal 1","Meal 2","Protein target","Daily kcal","Snack"], interactive=False)
            fig = gr.Plot()
            dl = gr.DownloadButton(label="Download plan (.md)")

            btn.click(
                predict_and_plan,
                inputs=[fasting_duration, meal_timing, weight, age, gender, height, activity, goal, diet, lang, ai_polish],
                outputs=[score_out, kpi_out, plan_md, plan_tbl, fig, dl],
                api_name="coach_plan"
            )

        # --- Tracker tab
        with gr.TabItem("Tracker"):
            with gr.Row():
                user = gr.Textbox(label="Username", value="")
                note = gr.Textbox(label="Note (optional)")
            with gr.Row():
                b1 = gr.Button("Start fast", variant="primary")
                b2 = gr.Button("End fast")
                b3 = gr.Button("Reload history")
            status = gr.Markdown("Not fasting.")
            hist = gr.Dataframe(interactive=False)
            hist_fig = gr.Plot()
            stats = gr.Markdown()

            b1.click(start_fast, inputs=[user, note], outputs=[status, note])
            b2.click(end_fast, inputs=[user], outputs=[status, hist, hist_fig, stats])
            b3.click(refresh_hist, inputs=[user], outputs=[hist, hist_fig, stats])
            demo.load(refresh_hist, inputs=[user], outputs=[hist, hist_fig, stats])

        # --- About tab
        with gr.TabItem("About"):
            gr.Markdown("""
            **How it works**  
            • Your predictor estimates a health score from inputs.  
            • The coach builds a 7‑day schedule matching your fasting window, goal, activity and diet style.  
            • Optional AI polish refines wording using a tiny model (ZeroGPU window).  
            • Tracker stores CSVs under `/data/` and never sends data elsewhere.
            """)

if __name__ == "__main__":
    demo.queue().launch()