import pickle
from flask import Flask, request, jsonify
from flask_cors import CORS
from sklearn.datasets import load_iris
from sklearn.ensemble import RandomForestClassifier

app = Flask(__name__)
CORS(app)

# --- Train or load model
try:
    model = pickle.load(open("model.pkl", "rb"))
except:
    iris = load_iris()
    X, y = iris.data, iris.target
    model = RandomForestClassifier()
    model.fit(X, y)
    pickle.dump(model, open("model.pkl", "wb"))

# --- Home route ---
@app.route("/", methods=["GET"])
def home():
    return """
    <h1>🌸 Iris Flower Classifier API</h1>
    <p>This is a Flask-based backend hosted on Hugging Face Spaces.</p>
    <p>It predicts the species of an Iris flower (<b>setosa</b>, <b>versicolor</b>, <b>virginica</b>) 
    based on sepal and petal measurements.</p>
    <p><b>Frontend UI:</b> 
    <a href="https://nielitropar.github.io/iris/" target="_blank">
        https://nielitropar.github.io/iris/
    </a></p>
    <p><b>API Endpoint:</b> <code>POST /predict</code></p>
    <p>Example JSON body:</p>
    <pre>{
  "sepal_length": 5.1,
  "sepal_width": 3.5,
  "petal_length": 1.4,
  "petal_width": 0.2
}</pre>
    <p>Response:</p>
    <pre>{
  "prediction": "setosa",
  "confidence": 0.98,
  "probabilities": {
      "setosa": 0.98,
      "versicolor": 0.01,
      "virginica": 0.01
  }
}</pre>
    """

# --- Predict route ---
@app.route("/predict", methods=["POST"])
def predict():
    data = request.json
    features = [
        data["sepal_length"],
        data["sepal_width"],
        data["petal_length"],
        data["petal_width"]
    ]

    # Predict class and probabilities
    prediction_idx = model.predict([features])[0]
    probs = model.predict_proba([features])[0]

    target_names = load_iris().target_names
    prediction_label = target_names[int(prediction_idx)]

    # Build probability dict
    probabilities = {
        target_names[i]: float(probs[i])
        for i in range(len(target_names))
    }

    confidence = float(max(probs))

    return jsonify({
        "prediction": prediction_label,
        "confidence": confidence,
        "probabilities": probabilities
    })


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