import joblib
import pandas as pd
import numpy as np
from flask import Flask, request, jsonify
from flask_cors import CORS

# Initialize Flask app
app = Flask("Pharmacy College Predictor")
CORS(app)

# Load trained model & helpers
pipeline = joblib.load('xgb_pipeline_gpu.pkl')
target_encoder = joblib.load('target_encoder.pkl')
choice_code_map = pd.read_csv('choice_code_map.csv').set_index('Choice Code')

# Home route
@app.get('/')
def home():
    return "✅ Welcome to Pharmacy College Predictor API!"

# Predict route
@app.post('/predict')
def predict():
    try:
        # Parse input JSON
        data = request.get_json()

        # Validate input
        required_fields = ['Category', 'Rank', 'Percentage']
        missing = [f for f in required_fields if f not in data]
        if missing:
            return jsonify({"error": f"Missing fields: {missing}"}), 400

        # Build DataFrame
        sample_df = pd.DataFrame([{
            'Category': data['Category'],
            'Rank': data['Rank'],
            'Percentage': data['Percentage']
        }])

        # Predict probabilities
        proba = pipeline.predict_proba(sample_df)[0]

        # Get top-20 indices (highest probabilities)
        top_20_idx = np.argsort(proba)[::-1][:20]

        # Normalize top-20 probs to sum to 100
        top_20_probs = proba[top_20_idx]
        top_20_probs_normalized = top_20_probs / top_20_probs.sum() * 100

        results = []
        for rank, (idx, prob) in enumerate(zip(top_20_idx, top_20_probs_normalized), start=1):
            choice_code = target_encoder.inverse_transform([idx])[0]
            row = choice_code_map.loc[int(choice_code)]
            college_name = row['College Name']
            results.append({
                "rank": rank,
                "choice_code": choice_code,
                "college_name": college_name,
                "probability_percent": round(float(prob), 2)
            })

        return jsonify({"top_20_predictions": results})

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

# Run server
if __name__ == '__main__':
    app.run(debug=False, host='0.0.0.0', port=7860)