from fastapi import FastAPI, UploadFile, File
from fastapi.responses import JSONResponse
from PIL import Image
import torch
import clip
import io

# -------------------------------
# Load CLIP
# -------------------------------
device = "cuda" if torch.cuda.is_available() else "cpu"
clip_model, preprocess = clip.load("ViT-B/32", device=device)

# -------------------------------
# Symbol Vocabulary (Math + Physics)
# -------------------------------
math_symbols = [
    "equal sign (=)", "plus sign (+)", "minus sign (-)", "division sign (/)",
    "multiplication sign (×)", "power or exponent (^)", "square root (√)",
    "integral (∫)", "summation (Σ)", "pi (π)", "infinity (∞)",
    "less than (<)", "greater than (>)", "less than or equal to (≤)",
    "greater than or equal to (≥)", "not equal to (≠)",
    "arrow (→)", "approximate (~)", "proportional (∝)"
]

greek_symbols = [
    "alpha (α)", "beta (β)", "gamma (γ)", "delta (δ)", "epsilon (ε)",
    "zeta (ζ)", "eta (η)", "theta (θ)", "lambda (λ)", "mu (μ)",
    "nu (ν)", "xi (ξ)", "rho (ρ)", "sigma (σ)", "tau (τ)",
    "phi (φ)", "chi (χ)", "psi (ψ)", "omega (ω)", "Omega (Ω)"
]

physics_symbols = [
    "Planck constant (ħ)", "speed of light (c)", "gravitational constant (G)",
    "Boltzmann constant (k)", "electric charge (e)", "permittivity (ε₀)",
    "permeability (μ₀)", "electron mass (mₑ)", "proton mass (mₚ)"
]

all_labels = math_symbols + greek_symbols + physics_symbols
text_tokens = clip.tokenize(all_labels).to(device)

# -------------------------------
# CLIP Detection Function
# -------------------------------
def detect_symbols(image: Image.Image, threshold: float = 0.15):
    image_tensor = preprocess(image).unsqueeze(0).to(device)
    with torch.no_grad():
        logits_per_image, _ = clip_model(image_tensor, text_tokens)
        probs = logits_per_image.softmax(dim=-1).cpu().numpy()[0]

    detected = []
    for label, p in zip(all_labels, probs):
        if p > threshold:  # filter noise
            detected.append({"symbol": label, "confidence": float(p)})

    return detected

# -------------------------------
# FastAPI App
# -------------------------------
app = FastAPI()

@app.post("/detect-symbols")
async def detect_math_physics_symbols(file: UploadFile = File(...)):
    try:
        contents = await file.read()
        image = Image.open(io.BytesIO(contents)).convert("RGB")

        results = detect_symbols(image)

        return JSONResponse(content={
            "filename": file.filename,
            "symbols_detected": results
        })
    except Exception as e:
        return JSONResponse(status_code=500, content={"error": str(e)})