from flask import Flask, request, jsonify
from flask_cors import CORS
import torch
import torch.nn as nn
from transformers import BertModel, BertTokenizer
import os

# ✅ Initialize Flask app
app = Flask(__name__)
CORS(app)  # Enable CORS for cross-origin requests

# ✅ Set device
device = "cuda" if torch.cuda.is_available() else "cpu"
print(f"Using device: {device}")

# ✅ Define the Correct Model
class YelpReviewClassifier(nn.Module):
    def __init__(self):
        super(YelpReviewClassifier, self).__init__()
        self.bert = BertModel.from_pretrained("bert-base-uncased")  # Load Pretrained BERT

        self.fc_text = nn.Linear(768, 128)  # Extracts text features
        self.fc_stars = nn.Linear(128, 5)   # Predicts star rating
        self.fc_useful = nn.Linear(128, 1)  # Predicts usefulness score

    def forward(self, input_ids, attention_mask):
        text_output = self.bert(input_ids=input_ids, attention_mask=attention_mask).pooler_output
        text_features = self.fc_text(text_output)

        stars_logits = self.fc_stars(text_features)  # Predict star ratings
        usefulness = self.fc_useful(text_features).squeeze()

        # Ensure usefulness is between 1 and 5
        usefulness = torch.round(torch.clamp(usefulness, 1, 5)).int()

        return stars_logits, usefulness

# ✅ Load Model
model_path = os.path.join(os.path.dirname(__file__), "models", "star_model.pth")
model = YelpReviewClassifier().to(device)

# ✅ Load weights
try:
    model.load_state_dict(torch.load(model_path, map_location=device), strict=False)
    model.eval()
    print("✅ Model loaded successfully!")
except Exception as e:
    print(f"❌ Error loading model: {e}")

# ✅ Load tokenizer
tokenizer = BertTokenizer.from_pretrained("bert-base-uncased")

# ✅ Define API Routes
@app.route("/")
def home():
    return jsonify({"message": "Welcome to Yelp Review AI Predictor!"})

@app.route("/predict", methods=["POST"])
def predict_rating():
    try:
        data = request.get_json()

        if not data:
            return jsonify({"error": "No input data provided"}), 400

        text = data.get("text", "").strip()
        if not text:
            return jsonify({"error": "Missing 'text' field"}), 400

        # ✅ Tokenize input
        tokens = tokenizer(
            text, truncation=True, padding="max_length", max_length=256, return_tensors="pt"
        )
        input_ids = tokens["input_ids"].to(device)
        attention_mask = tokens["attention_mask"].to(device)

        # ✅ Make prediction
        with torch.no_grad():
            stars_logits, usefulness = model(input_ids, attention_mask)
            predicted_stars = torch.argmax(stars_logits, dim=1).item() + 1  # Convert logits to 1-5 stars

        return jsonify({
            "predicted_stars": predicted_stars,
            "usefulness_score": usefulness.item()
        })

    except Exception as e:
        return jsonify({"error": str(e)}), 500

if __name__ == "__main__":
    app.run(host="0.0.0.0", port=8000, debug=True)  # Listen on all interfaces