from flask import Flask, render_template, request, jsonify
import joblib
import re

app = Flask(__name__)

# Load Model & Vectorizer
lr = joblib.load(open('tag_predictor_lr.pkl','rb'))
mlb_classes = joblib.load(open('mlb_classes.pkl','rb'))
tfidf = joblib.load(open('tfidf_vectorizer.pkl','rb'))

def preprocess_text(text):
    text = text.lower()
    text = re.sub(r'[^a-zA-Z\s]', '', text)
    text = re.sub(r'\s+', ' ', text).strip()
    return text

def predict_tags(text, model=lr, threshold=0.1):
    clean_text = preprocess_text(text)
    text_tfidf = tfidf.transform([clean_text])

    if hasattr(model, 'predict_proba'):
        probas = model.predict_proba(text_tfidf)
        tags = []
        p = []
        for i, class_name in enumerate(mlb_classes):
            if probas[i][0][1] > threshold:
                tags.append(class_name)
                p.append(probas[i][0][1])
        return tags, p
    else:
        preds = model.predict(text_tfidf)
        tags = [mlb_classes[i] for i, val in enumerate(preds[0]) if val == 1]
        return tags, preds[0]

@app.route('/')
def index():
    return render_template('index.html')

@app.route('/predict', methods=['POST'])
def predict():
    query = request.form.get("query")
    tags, probabilities = predict_tags(query)
    return jsonify({"tags": tags, "probabilities": probabilities})

if __name__ == '__main__':
    app.run(debug=True)