# -*- coding: utf-8 -*-
"""app.py

Automatically generated by Colab.

Original file is located at
    https://colab.research.google.com/drive/1E5d9dWFZwd3QoYrwG2SkR-slXGm2aMfL
"""

import pickle
import json
import torch
from fastapi import FastAPI
from pydantic import BaseModel
from transformers import BertTokenizer

# 1. Initialize App and Load Assets
# ==================================
app = FastAPI(
    title="TEXI'OR Sentiment Analysis API",
    description="An API to classify text into different sentiments using a fine-tuned BERT model.",
    version="1.0"
)

# Pydantic model for input data validation
class TextInput(BaseModel):
    text: str

# Use CUDA if available, otherwise CPU
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(f"✅ Using device: {device}")

# Load the fine-tuned model and tokenizer
try:
    with open('bert_sentiment_model (1).pkl', 'rb') as f:
        model = pickle.load(f)
    model.to(device)
    model.eval()  # Set model to evaluation mode
    print("✅ Model 'bert_sentiment_model (1).pkl' loaded successfully.")
except FileNotFoundError:
    print("❌ Model file not found. Make sure 'bert_sentiment_model (1).pkl' is uploaded.")
    model = None

try:
    with open('label_dict.json', 'r') as f:
        label_dict = json.load(f)
    # Create an inverse mapping from index to label string (e.g., {0: "nocode"})
    idx2label = {int(v): k for k, v in label_dict.items()}
    print("✅ Label dictionary 'label_dict.json' loaded successfully.")
except FileNotFoundError:
    print("❌ Label dictionary 'label_dict.json' not found.")
    idx2label = None

tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
print("✅ Tokenizer loaded successfully.")


# 2. Define API Endpoints
# ========================
@app.get("/")
def read_root():
    return {"message": "Welcome to the TEXI'OR API. Use the /predict endpoint to get a sentiment."}

@app.post("/predict")
def predict_sentiment(text_input: TextInput):
    """
    Predicts the sentiment of a given text.
    - Input: A JSON with a "text" field.
    - Output: A JSON with the predicted "sentiment" and "confidence" score.
    """
    if not all([model, idx2label, tokenizer]):
        return {"error": "API is not ready. A model, label, or tokenizer component failed to load."}

    text = text_input.text

    # Tokenize the input text
    inputs = tokenizer(
        text, padding=True, truncation=True, max_length=150, return_tensors='pt'
    )

    # Move tensors to the correct device
    input_ids = inputs['input_ids'].to(device)
    attention_mask = inputs['attention_mask'].to(device)

    # Get model predictions without calculating gradients
    with torch.no_grad():
        outputs = model(input_ids=input_ids, attention_mask=attention_mask)

    # Process the model output to get probabilities and the predicted class
    logits = outputs.logits
    probabilities = torch.nn.functional.softmax(logits, dim=-1)
    confidence, predicted_class_idx = torch.max(probabilities, dim=1)

    # Map the predicted index back to its string label
    predicted_label = idx2label.get(predicted_class_idx.item(), "Unknown Label")
    confidence_score = confidence.item()

    return {
        "sentiment": predicted_label,
        "confidence": round(confidence_score, 4)
    }