import numpy as np
import onnxruntime as ort
from PIL import Image
from transformers import AutoTokenizer

def predict_thumbnail_onnx(onnx_path: str, image_path: str, image_size: int = 299, threshold: float = 0.5):
    # 1. Load the ONNX model graph
    session = ort.InferenceSession(onnx_path)
    input_name = session.get_inputs()[0].name

    # 2. Preprocess the image (replicates T.Compose without PyTorch)
    img = Image.open(image_path).convert('RGB').resize((image_size, image_size))
    img_np = np.array(img, dtype=np.float32) / 255.0         # ToTensor part 1: Scale 0-1
    img_np = np.transpose(img_np, (2, 0, 1))                 # ToTensor part 2: HWC -> CHW

    # Normalize
    mean = np.array([0.485, 0.456, 0.406]).reshape(3, 1, 1)
    std = np.array([0.229, 0.224, 0.225]).reshape(3, 1, 1)
    img_np = (img_np - mean) / std

    # Add batch dimension (1, C, H, W)
    input_tensor = np.expand_dims(img_np, axis=0).astype(np.float32)

    # 3. Run inference through the ONNX graph
    outputs = session.run(None, {input_name: input_tensor})
    logit = outputs[0].item()

    # 4. Apply Sigmoid and threshold
    probability = 1 / (1 + np.exp(-logit))
    prediction = 1 if probability >= threshold else 0

    return probability, prediction

# --- Usage ---
#prob, pred = predict_thumbnail_onnx("./thumbnail_predictor.onnx", "/content/TBxS0XhdfmU-HD.jpg", image_size=299)
#print(f"Class: {pred}, Probability: {prob:.4f}")


def predict_title_onnx(onnx_path: str, title: str, tokenizer_name: str = "sentence-transformers/all-MiniLM-L6-v2", threshold: float = 0.5):
    # 1. Initialize the tokenizer
    # (In a production environment, load this once outside the function to save time)
    tokenizer = AutoTokenizer.from_pretrained(tokenizer_name)

    # 2. Tokenize the input text and return NumPy arrays ("np")
    inputs = tokenizer(
        title,
        return_tensors="np",
        padding="max_length",
        truncation=True,
        max_length=100 # 128 is plenty for YouTube titles
    )

    # Convert to int64, which is what ONNX exported transformer models expect
    input_ids = inputs["input_ids"].astype(np.int64)
    attention_mask = inputs["attention_mask"].astype(np.int64)

    # 3. Load the ONNX model graph
    session = ort.InferenceSession(onnx_path)

    # Prepare the input dictionary
    ort_inputs = {
        "input_ids": input_ids,
        "attention_mask": attention_mask
    }

    # Check if the exported ONNX graph also expects 'token_type_ids'
    # (MiniLM usually doesn't need them, but exporting to ONNX sometimes forces them in)
    expected_inputs = [inp.name for inp in session.get_inputs()]
    if "token_type_ids" in expected_inputs and "token_type_ids" in inputs:
        ort_inputs["token_type_ids"] = inputs["token_type_ids"].astype(np.int64)

    # 4. Run inference
    outputs = session.run(None, ort_inputs)

    # Outputs are shape (Batch, 1). Extract the raw logit.
    logit = outputs[0].item()

    # 5. Apply Sigmoid and threshold (because of BCEWithLogitsLoss)
    probability = 1 / (1 + np.exp(-logit))
    prediction = 1 if probability >= threshold else 0

    return probability, prediction
    
# --- Usage ---
#prob, pred = predict_title_onnx(onnx_path: str, title: str, tokenizer_name: str = "sentence-transformers/all-MiniLM-L6-v2", threshold: float = 0.5)
#print(f"Class: {pred}, Probability: {prob:.4f}")