from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from typing import List
from transformers import AutoTokenizer, AutoModelForSequenceClassification
import torch

# ── Config ────────────────────────────────────────────────────────────────────
HF_REPO    = 'ethnmcl/articulation-model'
MAX_LENGTH = 256
DEVICE     = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

SCORE_LABELS = {
    1: 'No Articulation',
    2: 'Minimal',
    3: 'Partial',
    4: 'Good',
    5: 'Full Articulation'
}

# ── Load model once at startup ─────────────────────────────────────────────────
print(f'Loading model from {HF_REPO}...')
tokenizer = AutoTokenizer.from_pretrained(HF_REPO)
model     = AutoModelForSequenceClassification.from_pretrained(HF_REPO)
model     = model.to(DEVICE)
model.eval()
print('Model ready.')

# ── App ───────────────────────────────────────────────────────────────────────
app = FastAPI(
    title='Articulation Scoring API',
    description='Scores how well a statement communicates technical progress on a 1–5 scale.',
    version='1.0.0'
)

# ── Schemas ───────────────────────────────────────────────────────────────────
class SingleRequest(BaseModel):
    statement: str

class BatchRequest(BaseModel):
    statements: List[str]

class ScoreResult(BaseModel):
    statement: str
    score: float
    rounded: int
    label: str

# ── Helpers ───────────────────────────────────────────────────────────────────
def run_inference(statements: List[str]) -> List[dict]:
    inputs = tokenizer(
        statements,
        return_tensors='pt',
        truncation=True,
        padding='max_length',
        max_length=MAX_LENGTH
    ).to(DEVICE)

    with torch.no_grad():
        preds_norm = model(**inputs).logits.squeeze(-1).cpu().tolist()

    if isinstance(preds_norm, float):
        preds_norm = [preds_norm]

    results = []
    for statement, pred in zip(statements, preds_norm):
        pred    = max(0.0, min(1.0, pred))
        score   = round(pred * 4 + 1, 2)
        rounded = round(score)
        results.append({
            'statement': statement,
            'score'    : score,
            'rounded'  : rounded,
            'label'    : SCORE_LABELS.get(rounded, 'Unknown')
        })
    return results

# ── Routes ────────────────────────────────────────────────────────────────────
@app.get('/')
def root():
    return {'status': 'ok', 'model': HF_REPO}

@app.post('/score', response_model=ScoreResult)
def score_single(req: SingleRequest):
    if not req.statement.strip():
        raise HTTPException(status_code=400, detail='Statement cannot be empty.')
    return run_inference([req.statement])[0]

@app.post('/score/batch', response_model=List[ScoreResult])
def score_batch(req: BatchRequest):
    if not req.statements:
        raise HTTPException(status_code=400, detail='Statements list cannot be empty.')
    if len(req.statements) > 50:
        raise HTTPException(status_code=400, detail='Max 50 statements per batch.')
    return run_inference(req.statements)