engine.py

"""DermaGlow engine: photo analysis, hybrid matching, product matching, weather advice.

Course constraints honored:
- Dataset is read directly from the HF dataset repo (havaferber/dermaglow-skin-faces).
- Embeddings + winning embedding model are read from the HF model repo (havaferber/dermaglow-embeddings,
  model = openai/clip-vit-base-patch32).
"""
import json
import numpy as np
import pandas as pd
import requests
import torch
from huggingface_hub import hf_hub_download
from transformers import AutoImageProcessor, AutoModel, AutoTokenizer

DATASET_REPO = "havaferber/dermaglow-skin-faces"
EMB_REPO = "havaferber/dermaglow-embeddings"
CLIP_ID = "openai/clip-vit-base-patch32"

SKIN_TYPES = ["oily", "dry", "combination", "normal"]
CONCERNS = ["acne", "redness", "pigmentation", "wrinkles", "dullness", "dehydration"]
GOAL_TO_CONCERN = {
    "Clear acne / breakouts": "acne",
    "Calm redness / irritation": "redness",
    "Fade dark spots / even tone": "pigmentation",
    "Anti-aging / smooth fine lines": "wrinkles",
    "Get my glow back": "dullness",
    "Deep hydration": "dehydration",
}


class Engine:
    def __init__(self):
        from datasets import load_dataset
        ds = load_dataset(DATASET_REPO, split="train")          # course: read from HF dataset repo
        self.meta = ds.remove_columns("image").to_pandas()

        emb_df = pd.read_parquet(hf_hub_download(EMB_REPO, "clip_embeddings.parquet"))
        order = {rid: i for i, rid in enumerate(self.meta["id"])}
        emb_df = emb_df.sort_values("id", key=lambda s: s.map(order))
        self.emb = np.array(emb_df["embedding"].tolist(), dtype=np.float32)   # (N, 512), L2-normalized

        self.clip_proc = AutoImageProcessor.from_pretrained(CLIP_ID)
        self.clip = AutoModel.from_pretrained(CLIP_ID).eval()
        self.clip_tok = AutoTokenizer.from_pretrained(CLIP_ID)

        import os
        self.catalog = pd.read_parquet(os.path.join(os.path.dirname(__file__), "sephora_catalog.parquet"))

    # ---------- photo analysis ----------
    def embed_photo(self, pil_img):
        with torch.no_grad():
            v = self.clip.get_image_features(**self.clip_proc(images=pil_img.convert("RGB"), return_tensors="pt"))
        v = v[0].numpy().astype(np.float32)
        return v / np.linalg.norm(v)

    def _zero_shot(self, pil_img, prompts):
        """CLIP zero-shot: returns softmax scores over the prompt dict {label: text}."""
        labels, texts = list(prompts.keys()), list(prompts.values())
        with torch.no_grad():
            ti = self.clip_tok(texts, padding=True, return_tensors="pt")
            ii = self.clip_proc(images=pil_img.convert("RGB"), return_tensors="pt")
            out = self.clip(**ti, **ii)
        probs = out.logits_per_image[0].softmax(dim=-1).numpy()
        return dict(zip(labels, probs))

    def analyze_photo(self, pil_img):
        """Estimate visible attributes from the photo with CLIP zero-shot."""
        tone = self._zero_shot(pil_img, {
            "I": "a face with very fair porcelain skin", "II": "a face with fair skin",
            "III": "a face with light-medium beige skin", "IV": "a face with olive tan skin",
            "V": "a face with brown skin", "VI": "a face with deep dark brown skin"})
        concern = self._zero_shot(pil_img, {
            "acne": "a face with acne breakouts and pimples",
            "redness": "a face with red irritated flushed skin",
            "pigmentation": "a face with dark spots and uneven pigmentation",
            "wrinkles": "a face with wrinkles and fine lines",
            "dullness": "a face with dull tired-looking skin",
            "dehydration": "a face with dry flaky dehydrated skin"})
        shine = self._zero_shot(pil_img, {
            "oily": "a face with oily shiny skin", "dry": "a face with dry rough skin",
            "normal": "a face with healthy balanced skin"})
        top = lambda d: max(d, key=d.get)
        return {"skin_tone": top(tone), "visible_concern": top(concern),
                "visible_type": top(shine), "concern_scores": concern}

    # ---------- hybrid recommendation ----------
    def match_profiles(self, photo_vec, skin_type, main_concern, age_group, k=3):
        """CLIP visual similarity + label boosts from the questionnaire (hybrid matcher)."""
        score = self.emb @ photo_vec
        score = score + 0.15 * (self.meta["skin_type"] == skin_type).values
        score = score + 0.15 * (self.meta["main_concern"] == main_concern).values
        score = score + 0.05 * (self.meta["age_group"] == age_group).values
        if "face_found" in self.meta:
            score = score - 0.5 * (~self.meta["face_found"].astype(bool)).values
        idx = np.argsort(-score)[:k]
        return self.meta.iloc[idx], score[idx]

    # ---------- product matching ----------
    STEP_CATEGORY = [
        (("sunscreen", "spf"), "Sunscreen"),
        (("cleanser", "wash", "cleansing", "micellar", "makeup remover"), "Cleansers"),
        (("eye",), "Eye Care"),
        (("mask",), "Masks"),
        (("moisturizer", "cream", "lotion", "balm"), "Moisturizers"),
        (("serum", "treatment", "toner", "essence", "exfoliant", "oil", "spot"), "Treatments"),
    ]
    # Friendly, singular label for each Sephora category — shown so the user knows what each product IS.
    CATEGORY_DISPLAY = {
        "Cleansers": "Cleanser", "Moisturizers": "Moisturizer", "Treatments": "Serum / treatment",
        "Sunscreen": "Sunscreen", "Eye Care": "Eye care", "Masks": "Mask",
    }

    def products_for_step(self, product_name, key_ingredients, n=2):
        pn = product_name.lower()
        category = next((cat for words, cat in self.STEP_CATEGORY if any(w in pn for w in words)), "Treatments")
        cand = self.catalog[self.catalog["secondary_category"] == category].copy()
        ingr_terms = [t for ing in key_ingredients for t in str(ing).lower().replace("%", " ").split() if len(t) > 3]
        cand["match"] = cand["ingredients"].apply(lambda s: sum(t in s for t in ingr_terms))
        cand["score"] = cand["match"] * 2 + cand["rating"].fillna(3.5) + np.log1p(cand["reviews"]) / 10
        cand["has_ingredient"] = cand["match"] > 0
        cand = cand.sort_values(["has_ingredient", "score"], ascending=False).head(n)
        recs = cand[["brand_name", "product_name", "price_usd", "rating", "url"]].to_dict("records")
        ptype = self.CATEGORY_DISPLAY.get(category, category)
        for r in recs:
            r["type"] = ptype
        return recs

    # ---------- live weather (bonus: real data) ----------
    @staticmethod
    def weather_advice(city):
        try:
            g = requests.get("https://geocoding-api.open-meteo.com/v1/search",
                             params={"name": city, "count": 1}, timeout=10).json()
            loc = g["results"][0]
            w = requests.get("https://api.open-meteo.com/v1/forecast",
                             params={"latitude": loc["latitude"], "longitude": loc["longitude"],
                                     "daily": "uv_index_max,temperature_2m_max,relative_humidity_2m_mean",
                                     "timezone": "auto"}, timeout=10).json()["daily"]
            uv, temp, hum = w["uv_index_max"][0], w["temperature_2m_max"][0], w["relative_humidity_2m_mean"][0]
        except Exception:
            return None
        tips = []
        if uv >= 8:
            tips.append(f"**Very high UV today ({uv})** — apply SPF 50 and reapply every 2 hours outdoors.")
        elif uv >= 5:
            tips.append(f"**High UV today ({uv})** — sunscreen is a must before leaving home.")
        elif uv >= 3:
            tips.append(f"**Moderate UV ({uv})** — your morning SPF step has you covered.")
        else:
            tips.append(f"**Low UV ({uv})** — SPF still recommended as a daily habit.")
        if temp >= 30:
            tips.append(f"**Hot day ({temp}°C)** — prefer light gel textures and rinse sweat off gently.")
        elif temp <= 8:
            tips.append(f"**Cold day ({temp}°C)** — use a richer moisturizer to protect your skin barrier.")
        if hum <= 35:
            tips.append(f"**Dry air ({hum}% humidity)** — add a hydrating toner or a few drops of hyaluronic serum.")
        elif hum >= 75:
            tips.append(f"**Humid air ({hum}% humidity)** — skip heavy creams tonight, let your skin breathe.")
        return {"city": loc["name"], "country": loc.get("country", ""), "uv": uv, "temp": temp,
                "humidity": hum, "tips": tips}