import torch
from transformers import PaliGemmaProcessor, PaliGemmaForConditionalGeneration

MODEL_PATH = "./paligemma_offline"

model = PaliGemmaForConditionalGeneration.from_pretrained(
    MODEL_PATH,
    torch_dtype=torch.bfloat16,
    device_map="auto",
    local_files_only=True
).eval()

processor = PaliGemmaProcessor.from_pretrained(
    MODEL_PATH,
    local_files_only=True
)

import cv2
import numpy as np
from PIL import Image

def preprocess_frame(frame_bgr: np.ndarray) -> Image.Image:
    # Resize to inference resolution before any VLM processing
    frame_small = cv2.resize(frame_bgr, (448, 448))
    # Convert BGR (OpenCV) to RGB (PIL/transformers)
    return Image.fromarray(cv2.cvtColor(frame_small, cv2.COLOR_BGR2RGB))

# A generic set of states tailored to your specific domain
VALID_CLASSES = ['active', 'idle', 'error', 'offline', 'unknown']

def classify_frame(model: torch.nn.Module, processor, image: Image.Image) -> str:
    prompt = (
        f"<image> What is the current operational state shown in this frame? "
        f"You MUST choose ONLY ONE from: {VALID_CLASSES}."
    )

    inputs = processor(
        text=prompt,
        images=image,
        return_tensors="pt"
    ).to(model.device, model.dtype)

    with torch.inference_mode():
        output_ids = model.generate(
            **inputs,
            max_new_tokens=10,   # A single class label needs at most 2-3 tokens
            do_sample=False      # Greedy decoding: deterministic, faster, no temperature needed
        )

    raw_output = processor.decode(
        output_ids[0][inputs.input_ids.shape[-1]:],
        skip_special_tokens=True
    ).strip().lower()

    # Sanitize to alphabetic characters only
    return ''.join(filter(str.isalpha, raw_output))

from collections import deque, Counter

class TemporalSmoother:
    def __init__(self, window_size: int = 5):
        self.history = deque(maxlen=window_size)

    def update(self, prediction: str) -> str:
        self.history.append(prediction)
        # Return the most common prediction in the window
        return Counter(self.history).most_common(1)[0][0]

import threading
import time
import numpy as np
import cv2
import torch
from typing import Optional, Any

class SharedState:
    """
    Thread-safe state container.
    The lock is strictly granular: it is only held for memory assignment/copying,
    never during expensive I/O or AI inference operations.
    """
    def __init__(self):
        self.latest_frame: Optional[np.ndarray] = None
        self.prediction: str = "WAITING"
        self.lock = threading.Lock()
        self.running: bool = True

shared = SharedState()

def video_capture_worker(source: int = 0) -> None:
    """Reads frames at hardware speed and updates the shared state."""
    cap = cv2.VideoCapture(source)

    while shared.running:
        ret, frame = cap.read()
        if not ret:
            time.sleep(0.01)
            continue

        with shared.lock:
            # Overwrite with the freshest data.
            # Pointer assignment is fast enough to barely hold the lock.
            shared.latest_frame = frame

    cap.release()

def inference_worker(model: torch.nn.Module, processor: Any) -> None:
    """Consumes the latest frame at the AI's maximum throughput rate."""
    smoother = TemporalSmoother(window_size=5)

    while shared.running:
        with shared.lock:
            # Deep copy to prevent OpenCV from mutating the array during inference
            frame = shared.latest_frame.copy() if shared.latest_frame is not None else None

        if frame is None:
            time.sleep(0.05)
            continue

        try:
            image = preprocess_frame(frame)
            raw_pred = classify_frame(model, processor, image)
            smoothed_pred = smoother.update(raw_pred)

            with shared.lock:
                shared.prediction = smoothed_pred

        except torch.cuda.OutOfMemoryError:
            # Handle temporary VRAM spikes gracefully without killing the thread
            with shared.lock:
                shared.prediction = "OOM_ERROR"
            time.sleep(1.0)

        except Exception as e:
            # Catch corrupt frames or tensor mismatches
            with shared.lock:
                shared.prediction = "ERROR"

# Initialize background workers as Daemon threads
threads = [
    threading.Thread(target=video_capture_worker, args=(0,), daemon=True),
    threading.Thread(target=inference_worker, args=(model, processor), daemon=True),
]

for t in threads:
    t.start()

# Main Thread UI Loop
# UI libraries (cv2.imshow) must run on the main thread to prevent OS-level crashes.
while shared.running:
    with shared.lock:
        frame = shared.latest_frame.copy() if shared.latest_frame is not None else None
        label = shared.prediction

    if frame is not None:
        cv2.putText(frame, label.upper(), (30, 50),
                    cv2.FONT_HERSHEY_SIMPLEX, 1.5, (0, 255, 0), 3)
        cv2.imshow("Live Classification", frame)

    # 30 FPS rendering limit (33ms) + graceful shutdown
    if cv2.waitKey(33) & 0xFF == ord('q'):
        shared.running = False

cv2.destroyAllWindows()