import pickle
import numpy as np
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from typing import Dict, List
import uvicorn

# Initialize FastAPI app
app = FastAPI(
    title="Sentiment Analysis API",
    description="A sentiment analysis model that predicts sentiment (positive/negative) from text",
    version="1.0.0"
)

# Load the model
try:
    with open("sentiment.pkl", "rb") as f:
        model = pickle.load(f)
except FileNotFoundError:
    raise Exception("sentiment.pkl not found. Please ensure the model file is in the correct location.")

# Pydantic models for request/response
class TextInput(BaseModel):
    text: str

class PredictionResponse(BaseModel):
    prediction: int
    confidence: float
    sentiment: str

class ProbabilityResponse(BaseModel):
    probabilities: List[float]
    prediction: int
    sentiment: str

@app.get("/")
async def root():
    return {
        "message": "Sentiment Analysis API",
        "endpoints": {
            "/predict": "Get sentiment prediction (0 or 1)",
            "/predict_proba": "Get prediction probabilities",
            "/health": "Check API health"
        }
    }

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

@app.post("/predict", response_model=PredictionResponse)
async def predict_sentiment(input_data: TextInput):
    """
    Predict sentiment from text input.
    Returns integer prediction (0 for negative, 1 for positive).
    """
    try:
        # Convert text to the format expected by your model
        # Note: You may need to preprocess the text depending on your model's requirements
        text_data = [input_data.text]
        
        # Get prediction
        prediction = model.predict(text_data)[0]
        
        # Get probabilities to calculate confidence
        probabilities = model.predict_proba(text_data)[0]
        confidence = float(max(probabilities))
        
        # Convert prediction to sentiment label
        sentiment = "positive" if prediction == 1 else "negative"
        
        return PredictionResponse(
            prediction=int(prediction),
            confidence=confidence,
            sentiment=sentiment
        )
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Prediction failed: {str(e)}")

@app.post("/predict_proba", response_model=ProbabilityResponse)
async def predict_probabilities(input_data: TextInput):
    """
    Get prediction probabilities for sentiment analysis.
    Returns probability arrays with shape (n_samples, 2).
    """
    try:
        # Convert text to the format expected by your model
        text_data = [input_data.text]
        
        # Get probabilities
        probabilities = model.predict_proba(text_data)[0]
        prediction = model.predict(text_data)[0]
        
        # Convert prediction to sentiment label
        sentiment = "positive" if prediction == 1 else "negative"
        
        return ProbabilityResponse(
            probabilities=probabilities.tolist(),
            prediction=int(prediction),
            sentiment=sentiment
        )
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Probability prediction failed: {str(e)}")

@app.post("/batch_predict")
async def batch_predict(texts: List[str]):
    """
    Predict sentiment for multiple texts at once.
    """
    try:
        predictions = model.predict(texts)
        probabilities = model.predict_proba(texts)
        
        results = []
        for i, text in enumerate(texts):
            sentiment = "positive" if predictions[i] == 1 else "negative"
            results.append({
                "text": text,
                "prediction": int(predictions[i]),
                "probabilities": probabilities[i].tolist(),
                "sentiment": sentiment,
                "confidence": float(max(probabilities[i]))
            })
        
        return {"results": results}
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Batch prediction failed: {str(e)}")

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