"""
NLP Evaluation API - Production Deployment
==========================================
Multi-task NLP model for interview evaluation
"""

from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
import torch
import torch.nn as nn
from transformers import RobertaTokenizer, RobertaModel
import uvicorn

app = FastAPI(title="NLP Evaluation API", version="1.0.0")

# Model configuration
TASKS = {
    "sentiment": {"labels": ["negative", "neutral", "positive"], "num_labels": 3},
    "emotion": {"labels": ["nervous", "confident", "stressed", "calm", "motivated"], "num_labels": 5},
    "communication": {"labels": ["poor", "fair", "good", "excellent"], "num_labels": 4},
    "confidence_level": {"labels": ["low", "medium", "high"], "num_labels": 3},
    "stress_level": {"labels": ["low", "medium", "high"], "num_labels": 3}
}

class MultiTaskRoBERTa(nn.Module):
    def __init__(self):
        super().__init__()
        self.roberta = RobertaModel.from_pretrained("roberta-base")
        hidden_size = 768
        self.classifiers = nn.ModuleDict({
            task: nn.Sequential(
                nn.Dropout(0.1),
                nn.Linear(hidden_size, hidden_size//2),
                nn.ReLU(),
                nn.Dropout(0.1),
                nn.Linear(hidden_size//2, info["num_labels"])
            )
            for task, info in TASKS.items()
        })
    
    def forward(self, input_ids, attention_mask):
        outputs = self.roberta(input_ids=input_ids, attention_mask=attention_mask)
        pooled = outputs.last_hidden_state[:, 0, :]
        return {task: classifier(pooled) for task, classifier in self.classifiers.items()}

# Load model
print("🧠 Loading model...")
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = MultiTaskRoBERTa().to(device)
checkpoint = torch.load("best_model_ultimate.pt", map_location=device)
model.load_state_dict(checkpoint["model_state_dict"])
model.eval()
print(f"✅ Model loaded on {device}")

# Load tokenizer
print("🔤 Loading tokenizer...")
tokenizer = RobertaTokenizer.from_pretrained("roberta-base")
print("✅ Tokenizer loaded")

class AnalyzeRequest(BaseModel):
    text: str

class AnalyzeResponse(BaseModel):
    sentiment: str
    emotion: str
    communication: str
    confidence_level: str
    stress_level: str
    scores: dict

@app.get("/")
def root():
    return {
        "message": "NLP Evaluation API",
        "version": "1.0.0",
        "status": "running",
        "model": "RoBERTa-base Multi-Task (126M params)",
        "tasks": list(TASKS.keys())
    }

@app.get("/health")
def health():
    return {"status": "healthy", "model_loaded": True}

@app.post("/analyze", response_model=AnalyzeResponse)
def analyze(request: AnalyzeRequest):
    try:
        # Tokenize
        encoding = tokenizer(
            request.text,
            max_length=256,
            padding="max_length",
            truncation=True,
            return_tensors="pt"
        )
        
        input_ids = encoding["input_ids"].to(device)
        attention_mask = encoding["attention_mask"].to(device)
        
        # Predict
        with torch.no_grad():
            logits = model(input_ids, attention_mask)
        
        # Get predictions and scores
        predictions = {}
        scores = {}
        
        for task, task_logits in logits.items():
            probs = torch.softmax(task_logits, dim=1)
            pred_idx = torch.argmax(probs, dim=1).item()
            pred_label = TASKS[task]["labels"][pred_idx]
            pred_score = probs[0][pred_idx].item()
            
            predictions[task] = pred_label
            scores[task] = round(pred_score, 4)
        
        return AnalyzeResponse(
            sentiment=predictions["sentiment"],
            emotion=predictions["emotion"],
            communication=predictions["communication"],
            confidence_level=predictions["confidence_level"],
            stress_level=predictions["stress_level"],
            scores=scores
        )
    
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

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