from flask import Flask, request, render_template, send_file, jsonify
import joblib
import pandas as pd
import numpy as np
import io

app = Flask(__name__)

# Load model and scaler
model = joblib.load("heart_model.pkl")
scaler = joblib.load("scaler.pkl")

FEATURES = ['age','sex','cp','trestbps','chol','fbs','restecg',
            'thalach','exang','oldpeak','slope','ca','thal']

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

# JSON endpoint for dynamic prediction
@app.route("/predict_json", methods=["POST"])
def predict_json():
    try:
        data_dict = request.get_json()
        data_df = pd.DataFrame([data_dict], columns=FEATURES)
        data_scaled = scaler.transform(data_df)
        pred = model.predict(data_scaled)[0]
        result = "No Heart Disease ✅" if pred == 0 else "Heart Disease ❌"
        return jsonify({"prediction": result})
    except Exception as e:
        return jsonify({"error": str(e)})

# Batch CSV prediction
@app.route("/batch_predict", methods=["POST"])
def batch_predict():
    try:
        file = request.files['file']
        if not file:
            return render_template("index.html", result="No file uploaded")

        df = pd.read_csv(file)
        if not all(col in df.columns for col in FEATURES):
            return render_template("index.html", result="CSV must have all required columns!")

        X_scaled = scaler.transform(df[FEATURES])
        preds = model.predict(X_scaled)
        df["prediction"] = ["No Heart Disease ✅" if p==0 else "Heart Disease ❌" for p in preds]

        # Return CSV as download
        output = io.StringIO()
        df.to_csv(output, index=False)
        output.seek(0)
        return send_file(io.BytesIO(output.getvalue().encode()),
                         mimetype="text/csv",
                         as_attachment=True,
                         download_name="batch_predictions.csv")

    except Exception as e:
        return render_template("index.html", result=f"Error: {str(e)}")

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