import tensorflow as tf
import numpy as np
from tensorflow.keras.models import load_model
from tensorflow.keras.preprocessing.text import tokenizer_from_json
from tensorflow.keras.preprocessing.sequence import pad_sequences
import json
import pickle
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel

app = FastAPI(title="News Source Classifier")

try:
    model = load_model('news_classifier.h5')
    
    with open('tokenizer.json') as f:
        tokenizer_data = json.load(f)
        tokenizer = tokenizer_from_json(tokenizer_data)
        
    with open('vectorizer.pkl', 'rb') as f:
        vectorizer = pickle.load(f)
except Exception as e:
    print(f"Error loading model: {str(e)}")
    raise

class PredictionRequest(BaseModel):
    text: str

class PredictionResponse(BaseModel):
    source: str
    confidence: float

@app.post("/predict", response_model=PredictionResponse)
async def predict(request: PredictionRequest):
    try:
        sequence = tokenizer.texts_to_sequences([request.text])
        padded = pad_sequences(sequence, maxlen=100)

        prediction = model.predict(padded)
        confidence = float(np.max(prediction))
        
        predicted_class = int(np.argmax(prediction))
        source = 'foxnews' if predicted_class == 0 else 'nbc'
        
        return PredictionResponse(
            source=source,
            confidence=confidence
        )
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.get("/")
async def root():
    return {
        "message": "News Source Classifier API",
        "usage": "Make a POST request to /predict with a JSON payload containing 'text' field"
    }