server.py

import os
import time
import csv
import json
import math
import base64
import datetime
import threading
import collections
import numpy as np
from PIL import Image, ImageDraw, ImageFont, ImageFilter
import io
import torch
import torch.nn.functional as F
from transformers import AutoImageProcessor, AutoModelForImageClassification
from flask import Flask, request, jsonify, render_template, Response
import requests as http_requests

# Optional imports — graceful fallback if not installed
try:
    import cv2
    CV2_AVAILABLE = True
except ImportError:
    CV2_AVAILABLE = False
    print("[WARN] cv2 not available — CLAHE disabled. Run: pip install opencv-python")

try:
    from scipy.ndimage import convolve as scipy_convolve
    SCIPY_AVAILABLE = True
except ImportError:
    SCIPY_AVAILABLE = False
    print("[WARN] scipy not available — Wiener deconvolution disabled.")


app = Flask(__name__)

# =============================================================================
#  PATHS
# =============================================================================
BASE_DIR        = os.path.dirname(os.path.abspath(__file__))
LOGS_DIR        = os.path.join(BASE_DIR, "logs")
CONFIG_PATH     = os.path.join(BASE_DIR, "config", "category_mapping.json")
ROI_PATH        = os.path.join(BASE_DIR, "config", "roi_config.json")
BACKGROUND_PATH = os.path.join(BASE_DIR, "config", "background_ref.npy")
MEMORY_PATH     = os.path.join(BASE_DIR, "config", "correction_memory.json")
CSV_LOG_PATH    = os.path.join(LOGS_DIR, "classification_log.csv")

os.makedirs(LOGS_DIR,                        exist_ok=True)
os.makedirs(os.path.join(BASE_DIR, "config"), exist_ok=True)


# =============================================================================
#  GLOBAL STATE
# =============================================================================
model                = None
processor            = None
category_mapping     = {}
last_processed_image = None
device               = "cpu"

# ── Frame pipeline ──────────────────────────────────────────────────────────
latest_raw_frame     = None
latest_annotated     = None
frame_lock           = threading.Lock()

# ── Background reference ──────────────────────────────────────────────────────
background_gray    = None
background_lock    = threading.Lock()
background_ready   = False

# ── Detection state ───────────────────────────────────────────────────────────
reference_gray     = None
reference_lock     = threading.Lock()
detect_confirm     = 0
PIXEL_DIFF_THRESH  = 25
CHANGED_PCT        = 0.08
CONFIRM_FRAMES     = 3
ADAPT_RATE         = 0.04
settle_until       = 0
SETTLE_DELAY_S     = 1.5

# ── Command queue ─────────────────────────────────────────────────────────────
command_queue      = []
command_lock       = threading.Lock()

# ── Classification result ─────────────────────────────────────────────────────
last_category      = None
last_category_time = 0
LABEL_SHOW_S       = 5.0
cycle_running      = False

# ── Gemma agent state ─────────────────────────────────────────────────────────
last_gemma_result  = {
    "status":            "idle",
    "prediction":        None,
    "timestamp":         None,
    "upcycling_tip":     None,
    "eco_fact":          None,
    "prep_instruction":  None,
    "witty_remark":      None,
    "material_analysis": None,
    "highlighted_type":  None,
    "highlighted_title": None,
    "highlighted_text":  None
}
gemma_lock         = threading.Lock()
gemma_busy         = False

# ── Memory database ───────────────────────────────────────────────────────────
correction_memory  = []
memory_lock        = threading.Lock()


# =============================================================================
#  ROI CONFIG
# =============================================================================
DEFAULT_ROI = {
    "cx":         160,
    "cy":          20,
    "angle_deg":   60,
    "radius_pct":  80,
    "zoom_factor": 1.0,
}
roi_config = dict(DEFAULT_ROI)


def load_roi():
    global roi_config
    if os.path.exists(ROI_PATH):
        try:
            with open(ROI_PATH) as f:
                roi_config = {**DEFAULT_ROI, **json.load(f)}
        except Exception:
            roi_config = dict(DEFAULT_ROI)

def save_roi():
    with open(ROI_PATH, "w") as f:
        json.dump(roi_config, f, indent=2)

def roi_radius():
    return int(240 * roi_config["radius_pct"] / 100)

def build_slice_mask(h, w):
    cx         = roi_config["cx"]
    cy         = roi_config["cy"]
    half_angle = roi_config["angle_deg"] / 2.0
    r          = roi_radius()
    ys, xs = np.mgrid[0:h, 0:w]
    dx = xs - cx
    dy = ys - cy
    dist  = np.sqrt(dx*dx + dy*dy)
    angle = np.abs(np.degrees(np.arctan2(dx, dy)))
    return (dist <= r) & (dist >= 1) & (angle <= half_angle)

def roi_bounding_box():
    cx       = roi_config["cx"]
    cy       = roi_config["cy"]
    r        = roi_radius()
    half_rad = math.radians(roi_config["angle_deg"] / 2.0)
    x0 = max(0,   int(cx - r * math.sin(half_rad)) - 4)
    y0 = max(0,   cy - 4)
    x1 = min(320, int(cx + r * math.sin(half_rad)) + 4)
    y1 = min(240, int(cy + r * math.cos(half_rad)) + 4)
    return x0, y0, x1, y1


# =============================================================================
#  BACKGROUND REFERENCE
# =============================================================================
def load_background():
    global background_gray, background_ready
    if os.path.exists(BACKGROUND_PATH):
        try:
            background_gray  = np.load(BACKGROUND_PATH).astype(np.float32)
            background_ready = True
        except Exception as e:
            print(f"[BG] Failed to load background: {e}")

def save_background(jpeg_bytes: bytes):
    global background_gray, background_ready
    img  = Image.open(io.BytesIO(jpeg_bytes)).convert("L").resize((320, 240))
    gray = np.array(img, dtype=np.float32)
    with background_lock:
        background_gray  = gray
        background_ready = True
    np.save(BACKGROUND_PATH, gray)

def apply_background_subtraction(img_np: np.ndarray, mask: np.ndarray) -> np.ndarray:
    with background_lock:
        bg = background_gray.copy() if background_ready else None
    if bg is None:
        return img_np
    gray = np.mean(img_np, axis=2).astype(np.float32)
    diff = np.abs(gray - bg)
    is_background = diff < 18.0
    result = img_np.copy()
    result[~mask] = 0
    result[is_background & mask] = 0
    return result


# =============================================================================
#  IMAGE PREPROCESSING PIPELINE
# =============================================================================
def apply_clahe(img: Image.Image) -> Image.Image:
    if not CV2_AVAILABLE:
        return img
    try:
        arr = np.array(img.convert("RGB"))
        lab = cv2.cvtColor(arr, cv2.COLOR_RGB2LAB)
        l, a, b = cv2.split(lab)
        clahe = cv2.createCLAHE(clipLimit=2.5, tileGridSize=(4, 4))
        l = clahe.apply(l)
        lab = cv2.merge([l, a, b])
        result = cv2.cvtColor(lab, cv2.COLOR_LAB2RGB)
        return Image.fromarray(result)
    except Exception:
        return img

def sharpen_defocus(img: Image.Image) -> Image.Image:
    try:
        img = img.filter(ImageFilter.UnsharpMask(radius=2.2, percent=175, threshold=2))
        img = img.filter(ImageFilter.EDGE_ENHANCE)
        return img
    except Exception:
        return img

def _fallback_bfs_component(mask: np.ndarray) -> np.ndarray:
    coords = np.argwhere(mask)
    if len(coords) == 0:
        return mask
    h, w = mask.shape
    visited = np.zeros((h, w), dtype=bool)
    largest_cluster = []
    for y, x in coords:
        if visited[y, x]:
            continue
        cluster = []
        queue = [(y, x)]
        visited[y, x] = True
        while queue:
            cy, cx = queue.pop(0)
            cluster.append((cy, cx))
            for dy in [-1, 0, 1]:
                for dx in [-1, 0, 1]:
                    if dy == 0 and dx == 0:
                        continue
                    ny, nx = cy + dy, cx + dx
                    if 0 <= ny < h and 0 <= nx < w:
                        if mask[ny, nx] and not visited[ny, nx]:
                            visited[ny, nx] = True
                            queue.append((ny, nx))
        if len(cluster) > len(largest_cluster):
            largest_cluster = cluster
    largest_mask = np.zeros((h, w), dtype=bool)
    for y, x in largest_cluster:
        largest_mask[y, x] = True
    return largest_mask

def preprocess_for_classification(jpeg_bytes: bytes) -> Image.Image:
    img = Image.open(io.BytesIO(jpeg_bytes)).convert("RGB").resize((320, 240))
    img_np = np.array(img)
    mask = build_slice_mask(240, 320)
    
    with background_lock:
        bg = background_gray.copy() if background_ready else None
        
    img_cropped = None
    if bg is not None:
        gray = np.mean(img_np, axis=2).astype(np.float32)
        diff = np.abs(gray - bg)
        raw_mask = (diff > 18.0) & mask
        
        # 1. Apply Median Filter to wipe out isolated discrete spots
        mask_pil = Image.fromarray((raw_mask * 255).astype(np.uint8))
        mask_pil = mask_pil.filter(ImageFilter.MedianFilter(size=5))
        filtered_mask = np.array(mask_pil) > 128
        
        # 2. Extract Largest Connected Component
        try:
            from scipy.ndimage import label
            labeled, num_features = label(filtered_mask)
            if num_features > 0:
                sizes = np.bincount(labeled.ravel())
                sizes[0] = 0 # Ignore background
                largest_label = np.argmax(sizes)
                final_mask = (labeled == largest_label)
            else:
                final_mask = filtered_mask
        except Exception:
            final_mask = _fallback_bfs_component(filtered_mask)
            
        object_coords = np.argwhere(final_mask)
        if len(object_coords) > 0:
            y_coords = object_coords[:, 0]
            x_coords = object_coords[:, 1]
            
            # Object center
            cy_obj = np.mean(y_coords)
            cx_obj = np.mean(x_coords)
            
            # Max distance to furthest continuous point
            dists = np.sqrt((x_coords - cx_obj)**2 + (y_coords - cy_obj)**2)
            r_max = np.max(dists)
            
            # Move boundary back by 20% to capture surrounding area (margin)
            half_side = max(15, int(r_max * 1.20))
            
            # Construct a perfect square side length
            side = 2 * half_side
            side = min(side, 240, 320)  # Keep within image bounds
            half_side = side // 2
            
            ymin = int(cy_obj - half_side)
            ymax = int(cy_obj + half_side)
            xmin = int(cx_obj - half_side)
            xmax = int(cx_obj + half_side)
            
            # Shift the window if it goes out of bounds to keep it square
            if xmin < 0:
                xmax -= xmin
                xmin = 0
            if xmax > 320:
                xmin -= (xmax - 320)
                xmax = 320
            if ymin < 0:
                ymax -= ymin
                ymin = 0
            if ymax > 240:
                ymin -= (ymax - 240)
                ymax = 240
            
            img_cropped = img.crop((xmin, ymin, xmax, ymax))
            print(f"[AI] Connected Object Square Crop: Center({cx_obj:.1f}, {cy_obj:.1f}), r_max={r_max:.1f}, box=({xmin},{ymin}) to ({xmax},{ymax})")
            
    if img_cropped is None:
        # Fallback to static ROI bounding box
        x0, y0, x1, y1 = roi_bounding_box()
        img_cropped = img.crop((x0, y0, x1, y1)) if (x1 > x0 and y1 > y0) else img
        print("[AI] Fallback to static ROI crop.")

    # In case the crop is not square (e.g. from static fallback), do center crop to square it
    w, h = img_cropped.size
    if w != h:
        min_dim = min(w, h)
        left = (w - min_dim) // 2
        top = (h - min_dim) // 2
        img_cropped = img_cropped.crop((left, top, left + min_dim, top + min_dim))
        print(f"[AI] Squared fallback crop: {w}x{h} -> {min_dim}x{min_dim}")

    # Resize directly and cleanly without aggressive sharpening filters to maintain natural quality
    img_cropped = img_cropped.resize((224, 224), Image.LANCZOS)
    return img_cropped


# =============================================================================
#  OVERLAY DRAWING
# =============================================================================
CATEGORY_COLORS = {
    "BIODEGRADABLE":     (0,   200,  0),
    "NON_BIODEGRADABLE": (220,  0,   0),
    "OTHER":             (220, 200,  0),
}
SLICE_COLOR = (255, 60, 60)
SLICE_ALPHA = 40

def draw_slice_overlay(img: Image.Image) -> Image.Image:
    cx         = roi_config["cx"]
    cy         = roi_config["cy"]
    half_angle = roi_config["angle_deg"] / 2.0
    r          = roi_radius()
    half_rad   = math.radians(half_angle)
    overlay = Image.new("RGBA", img.size, (0, 0, 0, 0))
    draw    = ImageDraw.Draw(overlay)
    steps  = 60
    points = [(cx, cy)]
    for s in range(steps + 1):
        a  = -half_rad + 2 * half_rad * s / steps
        ax = int(cx + r * math.sin(a))
        ay = int(cy + r * math.cos(a))
        points.append((ax, ay))
    points.append((cx, cy))
    draw.polygon(points, fill=(*SLICE_COLOR, SLICE_ALPHA))
    arc_lx = int(cx - r * math.sin(half_rad))
    arc_ly = int(cy + r * math.cos(half_rad))
    arc_rx = int(cx + r * math.sin(half_rad))
    draw.line([(cx, cy), (arc_lx, arc_ly)], fill=(*SLICE_COLOR, 230), width=2)
    draw.line([(cx, cy), (arc_rx, arc_ly)], fill=(*SLICE_COLOR, 230), width=2)
    for s in range(steps):
        a0 = -half_rad + 2 * half_rad * s       / steps
        a1 = -half_rad + 2 * half_rad * (s + 1) / steps
        p0 = (int(cx + r * math.sin(a0)), int(cy + r * math.cos(a0)))
        p1 = (int(cx + r * math.sin(a1)), int(cy + r * math.cos(a1)))
        draw.line([p0, p1], fill=(*SLICE_COLOR, 230), width=2)
    img = img.convert("RGBA")
    img = Image.alpha_composite(img, overlay)
    return img.convert("RGB")

def draw_category_tag(img: Image.Image, category: str) -> Image.Image:
    cx, cy = roi_config["cx"], roi_config["cy"]
    r = roi_radius()
    half_rad = math.radians(roi_config["angle_deg"] / 2.0)
    arc_lx = int(cx - r * math.sin(half_rad))
    arc_rx = int(cx + r * math.sin(half_rad))
    arc_by = int(cy + r * math.cos(half_rad))
    box_x0, box_y0 = max(0, arc_lx - 4), max(0, cy - 20)
    box_x1, box_y1 = min(319, arc_rx + 4), min(239, arc_by + 4)
    color = CATEGORY_COLORS.get(category, (180, 180, 0))
    draw  = ImageDraw.Draw(img)
    # Bounding box outline only (3px width)
    for t in range(3):
        draw.rectangle([(box_x0 + t, box_y0 + t), (box_x1 - t, box_y1 - t)], outline=color)
    # Top filled strip for label text
    strip_y1 = min(239, box_y0 + 14)
    draw.rectangle([(box_x0, box_y0), (box_x1, strip_y1)], fill=color)
    try:
        font = ImageFont.truetype("/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf", 10)
    except Exception:
        font = ImageFont.load_default()
    draw.text((box_x0 + 3, box_y0 + 2), category, fill=(0, 0, 0), font=font)
    return img

def annotate_frame(jpeg_bytes: bytes, category: str = None, show_overlays: bool = False) -> bytes:
    try:
        if not show_overlays:
            return jpeg_bytes
        img = Image.open(io.BytesIO(jpeg_bytes)).convert("RGB")
        img = draw_slice_overlay(img)
        if category:
            img = draw_category_tag(img, category)
        out = io.BytesIO()
        img.save(out, format="JPEG", quality=90)
        return out.getvalue()
    except Exception:
        return jpeg_bytes

def make_bg_preview(jpeg_bytes: bytes) -> bytes:
    try:
        img    = Image.open(io.BytesIO(jpeg_bytes)).convert("RGB")
        img_np = np.array(img)
        mask   = build_slice_mask(240, 320)
        with background_lock:
            bg = background_gray.copy() if background_ready else None
        if bg is not None:
            gray        = np.mean(img_np, axis=2).astype(np.float32)
            diff        = np.abs(gray - bg)
            is_bg_pixel = diff < 18.0
            overlay = img_np.copy()
            bg_mask = is_bg_pixel & mask
            overlay[bg_mask, 0] = np.clip(overlay[bg_mask, 0].astype(int) * 0.4, 0, 255).astype(np.uint8)
            overlay[bg_mask, 1] = np.clip(overlay[bg_mask, 1].astype(int) * 0.4 + 80, 0, 255).astype(np.uint8)
            overlay[bg_mask, 2] = np.clip(overlay[bg_mask, 2].astype(int) * 0.4, 0, 255).astype(np.uint8)
            outside = ~mask
            overlay[outside] = (overlay[outside].astype(int) * 0.25).astype(np.uint8)
            img = Image.fromarray(overlay)
        else:
            arr = img_np.copy()
            arr[~mask] = (arr[~mask].astype(int) * 0.3).astype(np.uint8)
            img = Image.fromarray(arr)
        img = draw_slice_overlay(img)
        out = io.BytesIO()
        img.save(out, format="JPEG", quality=80)
        return out.getvalue()
    except Exception:
        return jpeg_bytes


# =============================================================================
#  DETECTION LOOP
# =============================================================================
def frame_to_gray(jpeg_bytes: bytes) -> np.ndarray:
    img = Image.open(io.BytesIO(jpeg_bytes)).convert("L").resize((320, 240))
    return np.array(img, dtype=np.float32)

def update_reference(gray: np.ndarray):
    global reference_gray
    with reference_lock:
        if reference_gray is None: reference_gray = gray.copy()
        else: reference_gray = (1 - ADAPT_RATE) * reference_gray + ADAPT_RATE * gray

def check_for_waste(gray: np.ndarray) -> bool:
    global detect_confirm, reference_gray
    if time.time() < settle_until: return False
    with reference_lock:
        if reference_gray is None: return False
        ref = reference_gray.copy()
    mask = build_slice_mask(240, 320)
    diff = np.abs(gray - ref)
    changed = int(np.sum((diff > PIXEL_DIFF_THRESH) & mask))
    total   = int(np.sum(mask))
    trigger = int(total * CHANGED_PCT)
    if changed >= trigger:
        detect_confirm += 1
        if detect_confirm >= CONFIRM_FRAMES:
            detect_confirm = 0
            return True
    else:
        detect_confirm = 0
        update_reference(gray)
    return False

def detection_loop():
    global latest_annotated, last_category, last_category_time, cycle_running, settle_until
    while True:
        try:
            time.sleep(0.08)
            with frame_lock: raw = latest_raw_frame
            if raw is None: continue
            gray = frame_to_gray(raw)
            with reference_lock: have_ref = reference_gray is not None
            if not have_ref:
                update_reference(gray)
                continue
            
            # Check if there is active motion (changed pixels exceed threshold)
            motion_detected = False
            with reference_lock:
                if reference_gray is not None:
                    ref = reference_gray.copy()
                    mask = build_slice_mask(240, 320)
                    diff = np.abs(gray - ref)
                    changed = int(np.sum((diff > PIXEL_DIFF_THRESH) & mask))
                    total = int(np.sum(mask))
                    motion_detected = (changed >= int(total * CHANGED_PCT))
            
            cat = last_category if last_category and (time.time() - last_category_time) < LABEL_SHOW_S else None
            
            # Show overlays only when motion is detected, cycle is running, or recent category is active
            show_overlays = motion_detected or cycle_running or (cat is not None)
            
            annotated = annotate_frame(raw, cat, show_overlays=show_overlays)
            with frame_lock: latest_annotated = annotated
            
            if time.time() < settle_until: continue
            if check_for_waste(gray):
                result = classify_image(raw, log_to_csv=False)
                last_category = result["final_category"]
                last_category_time = time.time()
                settle_until = time.time() + 1.25
        except Exception: time.sleep(1.0)


def wait_for_queue_done():
    while True:
        time.sleep(0.1)
        with command_lock: empty = len(command_queue) == 0
        if empty:
            global cycle_running, reference_gray
            cycle_running = False
            with reference_lock: reference_gray = None
            break


# =============================================================================
#  GEMMA AGENT
# =============================================================================
OLLAMA_URL   = "http://localhost:11434/api/generate"
OLLAMA_MODEL = "gemma4"

STATIC_FALLBACKS = {
    "BIODEGRADABLE": {
        "upcycling_tip": "Compost this to nourish your garden soil or house plants.",
        "eco_fact": "Organic matter decomposes quickly, but in landfills without oxygen, it produces harmful methane gas.",
        "prep_instruction": "Remove any plastic packaging, stickers, or metal twist ties before composting.",
        "witty_remark": "Yum! Earthworms are going to love this delicious biodegradable snack.",
        "material_analysis": "Organic biodegradable waste (biological matter, cardboard, or clean paper)."
    },
    "NON_BIODEGRADABLE": {
        "upcycling_tip": "Clean and reuse plastic jars for storage, or turn metal cans into pencil holders.",
        "eco_fact": "A plastic bottle can take up to 450 years to decompose in a landfill.",
        "prep_instruction": "Empty all liquids, rinse out food residue, and flatten bottles or cans to save bin space.",
        "witty_remark": "Beep boop! This non-biodegradable item is going to stick around longer than my server uptime.",
        "material_analysis": "Recyclable or non-organic materials (plastics, metals, or glass)."
    },
    "OTHER": {
        "upcycling_tip": "Repurpose fabrics or old shoes into cleaning rags or workshop items.",
        "eco_fact": "Landfill space is finite, and mixed material items are the hardest to sort and recycle.",
        "prep_instruction": "Separate recyclable components from trash if possible. Do not put heavily contaminated items in recycling.",
        "witty_remark": "Sorting is hard, but together we make a great team! Let's handle this item responsibly.",
        "material_analysis": "General waste or mixed composition items (clothes, shoes, trash)."
    }
}

def _gemma_available() -> bool:
    try: return http_requests.get("http://localhost:11434/api/tags", timeout=2).status_code == 200
    except Exception: return False

def _image_to_b64(img: Image.Image) -> str:
    buf = io.BytesIO()
    img.save(buf, format="PNG")
    return base64.b64encode(buf.getvalue()).decode()

def text_to_robot_tones(text: str) -> str:
    words = text.split()
    tones = []
    import random
    rng = random.Random(hash(text))
    for word in words:
        syllables = max(1, len(word) // 3)
        for _ in range(syllables):
            freq = rng.randint(600, 1500)
            duration = rng.randint(45, 95)
            tones.append(f"{freq},{duration}")
    return "|".join(tones[:12])

def _run_gemma(img_crop: Image.Image, hf_result: dict) -> dict:
    cat = hf_result.get("final_category", "OTHER")
    fallback = STATIC_FALLBACKS.get(cat, STATIC_FALLBACKS["OTHER"])
    
    if not _gemma_available():
        res = {
            "status": "done",
            "prediction": cat,
            "timestamp": datetime.datetime.now().isoformat(),
            **fallback
        }
        return res

    try:
        img_b64 = _image_to_b64(img_crop)
        prompt = f"""
Analyze this waste item.
The primary image classifier has classified this item as: '{hf_result.get("detailed_label")}' (Category: {cat}).

You must respond ONLY with a single JSON object containing these exact keys:
- "upcycling_tip": A creative upcycling/reuse idea for this specific item.
- "eco_fact": An interesting environmental or decomposition fact about this material.
- "prep_instruction": How to prepare this item (e.g. rinse, remove caps, flatten, check for grease).
- "witty_remark": A short, witty, or encouraging comment in the persona of a friendly eco-conscious robot assistant.
- "material_analysis": A detailed breakdown of materials or recycle codes (e.g. PET 1, paper, aluminum).

Ensure the response is valid JSON and contains NO other text, no markdown block wrappers (like ```json), and no backticks.
"""
        payload = {"model": OLLAMA_MODEL, "prompt": prompt, "images": [img_b64], "stream": False}
        r = http_requests.post(OLLAMA_URL, json=payload, timeout=12)
        raw = r.json().get("response", "").strip()
        
        import re
        match = re.search(r'\{.*\}', raw, re.DOTALL)
        if match:
            json_data = json.loads(match.group(0))
        else:
            json_data = json.loads(raw)
            
        return {
            "status": "done",
            "prediction": cat,
            "timestamp": datetime.datetime.now().isoformat(),
            "upcycling_tip": json_data.get("upcycling_tip") or fallback["upcycling_tip"],
            "eco_fact": json_data.get("eco_fact") or fallback["eco_fact"],
            "prep_instruction": json_data.get("prep_instruction") or fallback["prep_instruction"],
            "witty_remark": json_data.get("witty_remark") or fallback["witty_remark"],
            "material_analysis": json_data.get("material_analysis") or fallback["material_analysis"]
        }
    except Exception as e:
        print(f"[Gemma Error] {e}")
        return {
            "status": "done",
            "prediction": cat,
            "timestamp": datetime.datetime.now().isoformat(),
            **fallback
        }

def gemma_analyze_async(img_crop: Image.Image, hf_result: dict, image_name: str):
    global gemma_busy, last_gemma_result
    cat = hf_result.get("final_category", "OTHER")
    with gemma_lock:
        last_gemma_result = {
            "status": "analyzing",
            "prediction": cat,
            "timestamp": datetime.datetime.now().isoformat(),
            "upcycling_tip": "...",
            "eco_fact": "...",
            "prep_instruction": "...",
            "witty_remark": "...",
            "material_analysis": "...",
            "highlighted_type": "...",
            "highlighted_title": "...",
            "highlighted_text": "..."
        }
    def _worker():
        global gemma_busy, last_gemma_result
        res = _run_gemma(img_crop, hf_result)
        
        import random
        insights = [
            ("upcycling", "💡 Upcycling Tip", res.get("upcycling_tip")),
            ("fact", "🌍 Eco-Fact", res.get("eco_fact")),
            ("prep", "⚙️ Prep Info", res.get("prep_instruction")),
            ("material", "🔍 Material Analysis", res.get("material_analysis")),
            ("witty", "🤖 Witty Remark", res.get("witty_remark")),
        ]
        valid_insights = [i for i in insights if i[2] and i[2] != "..."]
        if valid_insights:
            chosen_type, chosen_title, chosen_text = random.choice(valid_insights)
        else:
            chosen_type, chosen_title, chosen_text = ("witty", "🤖 Witty Remark", res.get("witty_remark") or "Sorting complete!")

        res["highlighted_type"] = chosen_type
        res["highlighted_title"] = chosen_title
        res["highlighted_text"] = chosen_text
        
        with gemma_lock:
            last_gemma_result = res
            
        tones = text_to_robot_tones(chosen_text)
        with command_lock:
            command_queue.append({"cmd": "ROBO_TALK", "param": tones})
            
        gemma_busy = False

    gemma_busy = True
    threading.Thread(target=_worker, daemon=True).start()


# =============================================================================
#  MEMORY DATABASE
# =============================================================================
def load_memory():
    global correction_memory
    if os.path.exists(MEMORY_PATH):
        try:
            with open(MEMORY_PATH) as f: correction_memory = json.load(f)
        except Exception: correction_memory = []

def save_correction(entry: dict):
    with memory_lock:
        correction_memory.append(entry)
        data = list(correction_memory)
    try:
        with open(MEMORY_PATH, "w") as f: json.dump(data, f, indent=2)
    except Exception: pass


# =============================================================================
#  AI CLASSIFICATION
# =============================================================================

def get_waste_attributes(label: str, category: str) -> dict:
    label_lower = label.lower()
    
    # Recyclable: plastic, cardboard, paper, metal, glass, cans, bottles
    recyclable_keywords = ["plastic", "cardboard", "paper", "metal", "glass", "can", "bottle", "scale"]
    is_recyclable = any(x in label_lower for x in recyclable_keywords)
    
    # Exclude non-recyclable items
    if label_lower in ["trash", "battery", "batteries", "electronics", "biological", "clothes", "shoes", "eraser"]:
        is_recyclable = False
        
    # Moisture: biological / organic are wet
    wet_keywords = ["biological", "organic", "food", "wet"]
    is_wet = any(x in label_lower for x in wet_keywords)
    
    return {
        "recyclable": "RECYCLABLE" if is_recyclable else "NON_RECYCLABLE",
        "moisture": "WET" if is_wet else "DRY"
    }

def classify_image(jpeg_bytes: bytes, log_to_csv: bool = True, bypass_preprocess: bool = False) -> dict:
    global last_processed_image
    try:
        if bypass_preprocess:
            img = Image.open(io.BytesIO(jpeg_bytes)).convert("RGB")
            # Center crop to a perfect square to prevent squeezing/stretching
            w, h = img.size
            if w != h:
                min_dim = min(w, h)
                left = (w - min_dim) // 2
                top = (h - min_dim) // 2
                img = img.crop((left, top, left + min_dim, top + min_dim))
            print(f"[AI] Bypassing classical CV crop for browser frame; center-cropped to square ({w}x{h} -> {img.size[0]}x{img.size[1]}).")
        else:
            img = preprocess_for_classification(jpeg_bytes)
        
        # Save preprocessed image for debug/viewing
        try:
            img.save(os.path.join(LOGS_DIR, "cropped_object.jpg"), format="JPEG", quality=90)
        except Exception:
            pass

        inputs = processor(images=img, return_tensors="pt")
        inputs = {k: v.to(device) for k, v in inputs.items()}
        with torch.no_grad():
            outputs = model(**inputs)
            probs   = torch.nn.functional.softmax(outputs.logits, dim=-1)[0]
        max_prob, max_idx = torch.max(probs, dim=-1)
        prob = max_prob.item() * 100
        raw_label = model.config.id2label[max_idx.item()]
        raw_label_lower = raw_label.lower()
        
        # Translate ResNet French labels to standard English mapping keywords
        french_to_english = {
            "batterie": "battery",
            "carton": "cardboard",
            "metal": "metal",
            "organique": "biological",
            "papier": "paper",
            "plastique": "plastic",
            "verre": "glass",
            "vetements": "clothes"
        }
        
        label = french_to_english.get(raw_label_lower, raw_label_lower)
        label_lower = label.lower()

        mapping_lower = {k.lower(): v for k, v in category_mapping.items()}
        if prob < 75.0 or label_lower not in mapping_lower:
            category = "OTHER"
            status   = "low_confidence" if prob < 75.0 else "unmapped_label"
        else:
            category = mapping_lower[label_lower]
            status   = "mapped"

        cid = f"waste_{datetime.datetime.now().strftime('%Y%m%d_%H%M%S')}"
        last_processed_image = cid
        attributes = get_waste_attributes(label, category)
        result = {
            "image_name":       cid,
            "detailed_label":   label,
            "model_confidence": round(prob, 2),
            "final_category":   category,
            "mapping_status":   status,
            "recyclable":       attributes["recyclable"],
            "moisture":         attributes["moisture"]
        }
        if log_to_csv:
            log_prediction_to_csv(datetime.datetime.now().isoformat(), cid, label, prob, category, status, "pending")
            gemma_analyze_async(img, result, cid)
        return result
    except Exception as e:
        print(f"[AI] Error: {e}")
        return {
            "image_name":       "error",
            "detailed_label":   "unknown",
            "model_confidence": 0.0,
            "final_category":   "OTHER",
            "mapping_status":   "error",
        }


# =============================================================================
#  CSV LOGGING
# =============================================================================
def log_prediction_to_csv(timestamp, image_name, detailed_label,
                           confidence, final_category, mapping_status, confirmation_status):
    file_exists = os.path.exists(CSV_LOG_PATH)
    os.makedirs(os.path.dirname(CSV_LOG_PATH), exist_ok=True)
    with open(CSV_LOG_PATH, mode="a", newline="", encoding="utf-8") as f:
        writer = csv.writer(f)
        if not file_exists:
            writer.writerow(["timestamp","image_name","detailed_label",
                             "model_confidence","final_category","mapping_status","confirmation_status"])
        writer.writerow([timestamp, image_name, detailed_label,
                         f"{confidence:.2f}", final_category, mapping_status, confirmation_status])

def update_csv_confirmation(image_name, new_status="confirmed"):
    if not os.path.exists(CSV_LOG_PATH): return False
    rows = []
    with open(CSV_LOG_PATH, mode="r", newline="", encoding="utf-8") as f:
        reader = csv.reader(f)
        for row in reader: rows.append(row)
    for row in rows:
        if len(row) > 1 and row[1] == image_name: row[6] = new_status
    with open(CSV_LOG_PATH, mode="w", newline="", encoding="utf-8") as f:
        csv.writer(f).writerows(rows)
    return True


# =============================================================================
#  COMMAND QUEUE
# =============================================================================
def enqueue_sort_cycle(category: str):
    esp_param = "OTHER"
    if category == "BIODEGRADABLE":
        esp_param = "BIO"
    elif category == "NON_BIODEGRADABLE":
        esp_param = "NONBIO"
    with command_lock:
        command_queue.clear()
        command_queue.extend([
            {"cmd": "SET_LED", "param": category},
            {"cmd": "BEEP"},
            {"cmd": "SERVO_A", "param": esp_param},
            {"cmd": "SERVO_B", "param": "OPEN"},
            {"cmd": "DELAY", "param": "750"},
            {"cmd": "SERVO_B", "param": "CLOSE"},
            {"cmd": "CLEAR_LEDS"},
            {"cmd": "SERVO_A", "param": "HOME"},
            {"cmd": "BEEP"}
        ])


# =============================================================================
#  STARTUP
# =============================================================================
def load_resources():
    global model, processor, category_mapping, device
    with open(CONFIG_PATH) as f: category_mapping = json.load(f)
    
    model_path = "dan-lara/Garbage-Classifier-Resnet-50-Finetuning"
    print(f"[AI] Loading model from Hugging Face Hub: {model_path}")
        
    device = "cuda" if torch.cuda.is_available() else "cpu"
    print(f"[AI] Running on device: {device}")
    

    
    # Workaround: Model repo has pytorch_model.pth instead of pytorch_model.bin
    from huggingface_hub import snapshot_download
    import shutil
    try:
        local_dir = snapshot_download(model_path)
        pth_path = os.path.join(local_dir, "pytorch_model.pth")
        bin_path = os.path.join(local_dir, "pytorch_model.bin")
        if os.path.exists(pth_path) and not os.path.exists(bin_path):
            print(f"[AI] Copying pytorch_model.pth to pytorch_model.bin: {bin_path}")
            shutil.copy(pth_path, bin_path)
        load_path = local_dir
    except Exception as e:
        print(f"[AI] Warning: Failed to download snapshot. Trying direct load fallback. Error: {e}")
        load_path = model_path
    
    processor  = AutoImageProcessor.from_pretrained(load_path)
    model      = AutoModelForImageClassification.from_pretrained(load_path).to(device)
    load_roi()
    load_background()
    load_memory()
    
    # Initialize empty CSV file with header if it doesn't exist
    if not os.path.exists(CSV_LOG_PATH):
        os.makedirs(os.path.dirname(CSV_LOG_PATH), exist_ok=True)
        with open(CSV_LOG_PATH, mode="w", newline="", encoding="utf-8") as f:
            csv.writer(f).writerow(["timestamp","image_name","detailed_label",
                                     "model_confidence","final_category","mapping_status","confirmation_status"])
                                     
    threading.Thread(target=detection_loop, daemon=True).start()


# =============================================================================
#  FLASK ROUTES
# =============================================================================
@app.route("/", methods=["GET"])
def index(): return render_template("index.html")

def apply_zoom(jpeg_bytes: bytes) -> bytes:
    zoom = roi_config.get("zoom_factor", 1.0)
    if zoom <= 1.0:
        return jpeg_bytes
    try:
        img = Image.open(io.BytesIO(jpeg_bytes))
        w, h = img.size
        crop_w = w / zoom
        crop_h = h / zoom
        x0 = (w - crop_w) / 2
        y0 = (h - crop_h) / 2
        x1 = x0 + crop_w
        y1 = y0 + crop_h
        img_cropped = img.crop((x0, y0, x1, y1)).resize((w, h), Image.LANCZOS)
        out = io.BytesIO()
        img_cropped.save(out, format="JPEG", quality=85)
        return out.getvalue()
    except Exception as e:
        print(f"[ZOOM] Preprocessing failed: {e}")
        return jpeg_bytes

@app.route("/api/upload_frame", methods=["POST"])
def upload_frame():
    global latest_raw_frame
    if not request.data:
        return jsonify({"error": "No data"}), 400
    zoomed = apply_zoom(request.data)
    with frame_lock: latest_raw_frame = zoomed
    return "", 200

@app.route("/api/live_feed", methods=["GET"])
def live_feed():
    with frame_lock: frame = latest_annotated or latest_raw_frame
    return Response(frame, mimetype="image/jpeg") if frame else ("", 404)

@app.route("/api/command", methods=["GET"])
def get_command():
    with command_lock: cmd = command_queue[0] if command_queue else {"cmd": "NONE"}
    return jsonify(cmd)

@app.route("/api/command_done", methods=["POST"])
def command_done():
    with command_lock:
        if command_queue:
            command_queue.pop(0)
            if not command_queue:
                threading.Thread(target=wait_for_queue_done, daemon=True).start()
    return "", 200

@app.route("/api/user_connected", methods=["POST"])
def user_connected():
    with command_lock: command_queue.append({"cmd": "PLAY_NOKIA", "param": ""})
    return jsonify({"status": "queued"})

@app.route("/api/roi", methods=["GET", "POST"])
def manage_roi():
    if request.method == "POST":
        data = request.get_json()
        for key in ("cx", "cy", "angle_deg", "radius_pct", "zoom_factor"):
            if key in data: roi_config[key] = float(data[key])
        save_roi()
    return jsonify(roi_config)

@app.route("/api/roi/reset", methods=["POST"])
def reset_roi():
    global roi_config
    roi_config = dict(DEFAULT_ROI)
    save_roi()
    return jsonify({"status": "ok", "roi": roi_config}), 200

@app.route("/api/capture_background", methods=["POST"])
def capture_background():
    with frame_lock: raw = latest_raw_frame
    if not raw: return jsonify({"error": "No frame"}), 400
    save_background(raw)
    return jsonify({"status": "ok"})

@app.route("/api/background_preview", methods=["GET"])
def background_preview():
    with frame_lock: raw = latest_raw_frame
    return Response(make_bg_preview(raw), mimetype="image/jpeg") if raw else ("", 404)


@app.route("/api/debug_crop", methods=["GET"])
def debug_crop():
    crop_path = os.path.join(LOGS_DIR, "cropped_object.jpg")
    if os.path.exists(crop_path):
        with open(crop_path, "rb") as f:
            data = f.read()
        return Response(data, mimetype="image/jpeg")
    return "No cropped image available yet.", 404

@app.route("/classify", methods=["POST"])
def classify_endpoint():
    global latest_raw_frame, last_category, last_category_time, settle_until, cycle_running
    
    # Check trigger_hardware parameter
    trigger_hardware = request.args.get("trigger_hardware", "true").lower() == "true"
    
    if trigger_hardware and cycle_running:
        return jsonify({"error": "Cycle already in progress"}), 409
    
    image_data = request.data or latest_raw_frame
    if not image_data:
        return jsonify({"error": "No image data"}), 400
        
    if trigger_hardware:
        cycle_running = True
        
    is_browser = bool(request.data)
    result = classify_image(image_data, log_to_csv=True, bypass_preprocess=is_browser)
    category = result.get("final_category")
    
    if category:
        if trigger_hardware:
            enqueue_sort_cycle(category)
            last_category = category
            last_category_time = time.time()
            settle_until = time.time() + SETTLE_DELAY_S + 4.0
        else:
            last_category = category
            last_category_time = time.time()
    else:
        if trigger_hardware:
            cycle_running = False
            
    return jsonify(result), 200

@app.route("/api/trigger_sort", methods=["POST"])
def trigger_sort():
    global latest_raw_frame, last_category, last_category_time, settle_until, cycle_running
    if cycle_running:
        return jsonify({"status": "ignored", "reason": "cycle_active"}), 200
    with frame_lock:
        img_data = latest_raw_frame
    if not img_data:
        return jsonify({"status": "ignored", "reason": "no_frame"}), 200
    cycle_running = True
    print("[TRIGGER] Hardware triggered sorting cycle.")
    result = classify_image(img_data, log_to_csv=True)
    category = result.get("final_category", "OTHER")
    enqueue_sort_cycle(category)
    last_category = category
    last_category_time = time.time()
    settle_until = time.time() + SETTLE_DELAY_S + 4.0
    return jsonify({"status": "triggered", "category": category, "result": result}), 200

@app.route("/api/last_gemma", methods=["GET"])
def get_last_gemma():
    with gemma_lock:
        res = dict(last_gemma_result)
    return jsonify(res), 200

@app.route("/api/feedback", methods=["POST"])
def feedback():
    data = request.get_json(silent=True) or {}
    with gemma_lock:
        gr = dict(last_gemma_result)

    entry = {
        "timestamp":          datetime.datetime.now().isoformat(),
        "image_name":         data.get("image_name", gr.get("image_name", "")),
        "upcycling_tip":      gr.get("upcycling_tip", ""),
        "eco_fact":           gr.get("eco_fact", ""),
        "prep_instruction":   gr.get("prep_instruction", ""),
        "witty_remark":       gr.get("witty_remark", ""),
        "material_analysis":  gr.get("material_analysis", ""),
        "hf_prediction":      gr.get("prediction", ""),
        "user_confirmed":     data.get("confirmed_category", ""),
        "was_correct":        data.get("was_correct", True),
    }
    save_correction(entry)

    # Also update CSV
    update_csv_confirmation(entry["image_name"], "confirmed")
    return jsonify({"status": "saved"}), 200


# ── System state ──────────────────────────────────────────────────────────────
@app.route("/api/is_active", methods=["GET"])
def is_active():
    return jsonify({"active": True}), 200

@app.route("/api/handshake", methods=["GET"])
def handshake():
    global cycle_running, reference_gray, command_queue
    with command_lock:
        command_queue.clear()
    with reference_lock:
        reference_gray = None
    cycle_running = False
    print("[HANDSHAKE] State synchronized.")
    return jsonify({"status": "synchronized"}), 200

@app.route("/confirm", methods=["POST"])
def confirm_classification():
    global last_processed_image
    data       = request.get_json(silent=True) or {}
    image_name = data.get("image_name") or last_processed_image
    if not image_name:
        return jsonify({"error": "No transaction to confirm"}), 400
    success = update_csv_confirmation(image_name)
    if success:
        return jsonify({"status": "success", "confirmed_image": image_name}), 200
    return jsonify({"error": "Not found"}), 404


# ── Categories ────────────────────────────────────────────────────────────────
@app.route("/api/categories", methods=["GET", "POST"])
def manage_categories():
    global category_mapping
    if request.method == "POST":
        try:
            new_mapping = request.get_json()
            if not new_mapping or not isinstance(new_mapping, dict):
                return jsonify({"error": "Invalid mapping"}), 400
            with open(CONFIG_PATH, "w") as f:
                json.dump(new_mapping, f, indent=2)
            category_mapping = new_mapping
            return jsonify({"status": "success"}), 200
        except Exception as e:
            return jsonify({"error": str(e)}), 500
    return jsonify(category_mapping), 200


# ── Logs ──────────────────────────────────────────────────────────────────────
@app.route("/api/logs", methods=["GET"])
def get_logs():
    if not os.path.exists(CSV_LOG_PATH):
        return jsonify([]), 200
    try:
        with open(CSV_LOG_PATH, mode="r", newline="", encoding="utf-8") as f:
            logs = list(csv.DictReader(f))
        return jsonify(logs[::-1][:20]), 200
    except Exception as e:
        return jsonify({"error": str(e)}), 500


# ── Background Status ─────────────────────────────────────────────────────────
@app.route("/api/background_status", methods=["GET"])
def background_status():
    return jsonify({"ready": background_ready}), 200


# =============================================================================
if __name__ == "__main__":
    load_resources()
    port = int(os.environ.get("PORT", 7860))
    app.run(host="0.0.0.0", port=port, debug=False, threaded=True)