import torch
import torch.nn.functional as F
from transformers import AutoTokenizer, AutoModelForSequenceClassification
import re
from fastapi import FastAPI
from pydantic import BaseModel

# ================== CONFIG ==================

MODEL_ID = "AkCh12/claim-extractor-deberta-v3"
MAX_LENGTH = 128
DEFAULT_THRESHOLD = 0.6

device = "cpu"

# ================== LOAD MODEL ==================

print("[INFO] Loading tokenizer...")
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)

print("[INFO] Loading model...")
model = AutoModelForSequenceClassification.from_pretrained(
    MODEL_ID,
    num_labels=2,
    torch_dtype=torch.float32
)

model.to(device)
model.eval()

torch.set_num_threads(1)

print("[INFO] Warming up model...")
with torch.no_grad():
    dummy = tokenizer("Warmup text.", return_tensors="pt")
    _ = model(**dummy)

print("[INFO] Model ready.")

# ================== HELPERS ==================

def split_sentences(text: str):
    return re.split(r'(?<=[.!?])\s+', text)


def score_sentences(sentences):
    inputs = tokenizer(
        sentences,
        padding=True,
        truncation=True,
        max_length=MAX_LENGTH,
        return_tensors="pt"
    )

    with torch.no_grad():
        logits = model(**inputs).logits
        probs = F.softmax(logits, dim=-1)

    return probs[:, 1].tolist()


def extract_claims(text: str, threshold: float):
    if not text or len(text.strip()) < 20:
        return []

    sentences = [s.strip() for s in split_sentences(text) if len(s.strip()) > 10]
    if not sentences:
        return []

    scores = score_sentences(sentences)

    return [
        {
            "sentence": s,
            "claim_probability": round(sc, 4),
            "is_claim": sc >= threshold
        }
        for s, sc in zip(sentences, scores)
    ]

# ================== FASTAPI ==================

app = FastAPI(title="Claim Extraction API")

class AnalyzeRequest(BaseModel):
    text: str
    threshold: float | None = None

@app.post("/extract")
def extract_api(payload: AnalyzeRequest):
    threshold = payload.threshold or DEFAULT_THRESHOLD

    claims = extract_claims(payload.text, threshold)

    return {
        "threshold": threshold,
        "num_sentences": len(claims),
        "claims": claims
    }

@app.get("/")
def home():
    return {
        "message": "Claim Extraction API running",
        "model": MODEL_ID
    }

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=7860)