# ============================================
# 💜 NasFit Vision AI — Dark Neon Edition
# ============================================

# 🚧 Asegúrate de tener este contenido en tu requirements.txt:
# transformers>=4.43.0
# torch
# accelerate
# gradio
# Pillow

import gradio as gr
import torch
from transformers import AutoModelForVision2Seq
from transformers import LlavaOneVisionProcessor  # 👈 clase correcta
from PIL import Image
import re

# ============================================
# 🔮 CARGA DEL MODELO LOCAL
# ============================================

MODEL_ID = "lmms-lab/llava-onevision-1.5-8b-instruct"

print("⏳ Cargando modelo local con trust_remote_code=True...")

processor = LlavaOneVisionProcessor.from_pretrained(MODEL_ID, trust_remote_code=True)
model = AutoModelForVision2Seq.from_pretrained(
    MODEL_ID,
    trust_remote_code=True,
    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
    device_map="auto"
)

print("✅ Modelo LLaVA-OneVision cargado correctamente.")

# ============================================
# 🧠 FUNCIONES DE ANÁLISIS NUTRICIONAL
# ============================================

def extract_macros(text):
    """Extrae proteínas, carbohidratos y grasas del texto generado."""
    def find_value(keyword):
        m = re.search(rf"{keyword}[^0-9]*([0-9]+)", text.lower())
        return int(m.group(1)) if m else 0
    p, c, f = find_value("prote"), find_value("carb"), find_value("gras")
    kcal = p * 4 + c * 4 + f * 9 if any([p, c, f]) else 0
    return {"protein": p, "carbs": c, "fat": f, "kcal": kcal}


def build_macro_card(macros):
    """Genera el HTML visual con barras de progreso tipo dashboard."""
    if not any(macros.values()):
        return "<div class='card'>⚖️ No se pudieron estimar los macros.</div>"

    def bar_html(value, color):
        width = min(value, 100)
        return f"""
        <div class='bar-bg'>
            <div class='bar-fill' style='width:{width}%; background:{color};'></div>
        </div>
        """

    return f"""
    <div class='card'>
      <h2>🍽️ Estimación Nutricional</h2>
      <div class='macro'><span>💪 Proteínas</span><span>{macros['protein']} g</span></div>
      {bar_html(macros['protein'], '#b25eff')}
      <div class='macro'><span>🥔 Carbohidratos</span><span>{macros['carbs']} g</span></div>
      {bar_html(macros['carbs'], '#00f0ff')}
      <div class='macro'><span>🥑 Grasas</span><span>{macros['fat']} g</span></div>
      {bar_html(macros['fat'], '#ff5efb')}
      <div class='macro kcal'><span>🔥 Calorías Totales</span><span>{macros['kcal']} kcal</span></div>
    </div>
    """


def analyze_food(image, text_prompt="Describe esta comida y estima sus calorías, proteínas, carbohidratos y grasas."):
    """Procesa la imagen localmente con el modelo y devuelve descripción + macros."""
    try:
        inputs = processor(text=text_prompt, images=image, return_tensors="pt").to(model.device)
        out = model.generate(**inputs, max_new_tokens=400)
        answer = processor.decode(out[0], skip_special_tokens=True)

        macros = extract_macros(answer)
        card = build_macro_card(macros)
        return f"<div class='desc'>{answer}</div>{card}"

    except Exception as e:
        return f"<div class='card error'>⚠️ Error: {e}</div>"

# ============================================
# 💅 INTERFAZ DE USUARIO (Purple Glassmorphism)
# ============================================

def build_interface():
    with gr.Blocks(css="""
/* --- DARK NEON THEME --- */
body {
  background: radial-gradient(circle at 20% 20%, #0d001f, #000);
  color: #fff;
  font-family: 'Inter', sans-serif;
}
.gradio-container {background: transparent !important;}
.card {
  backdrop-filter: blur(12px);
  background: rgba(30, 0, 60, 0.3);
  border: 1px solid rgba(200, 100, 255, 0.2);
  border-radius: 16px;
  padding: 1.2em;
  margin-top: 1em;
  box-shadow: 0 0 25px rgba(180, 0, 255, 0.15);
}
h1,h2 {color:#c18fff;}
.bar-bg {
  width:100%; height:8px; border-radius:6px;
  background:rgba(255,255,255,0.1); margin:4px 0 12px 0;
  overflow:hidden;
}
.bar-fill {height:100%; border-radius:6px; transition:width 1s ease;}
.macro {display:flex; justify-content:space-between; font-size:0.95em;}
.kcal {font-weight:600; color:#ffb3ff;}
.desc {
  background:rgba(255,255,255,0.05);
  padding:1em; border-radius:10px; line-height:1.5em;
  box-shadow:inset 0 0 20px rgba(180,0,255,0.1);
}
button {
  background:linear-gradient(90deg,#b25eff,#00f0ff);
  color:#fff; border:none; border-radius:12px;
  font-weight:600; transition:opacity .2s;
}
button:hover {opacity:0.8;}
""") as demo:

        gr.Markdown("""
        <h1>💜 NasFit Vision AI</h1>
        <p>Analiza tus comidas con IA y obtené tu ficha nutricional instantánea.</p>
        """)

        with gr.Row():
            with gr.Column(scale=1):
                img = gr.Image(label="📸 Imagen del plato", type="pil")
                txt = gr.Textbox(label="💬 Instrucción (opcional)", placeholder="Ej: ¿Cuántas calorías tiene este plato?")
                btn = gr.Button("🔍 Analizar", variant="primary")
            with gr.Column(scale=1):
                out = gr.HTML(label="🧠 Resultado")

        btn.click(analyze_food, [img, txt], out)

    return demo


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