from flask import Flask, render_template, request, url_for
import pickle
import numpy as np
import os
from ocr_utils import extract_keywords_from_report, score_text_for_risk 
from image_utils import predict_xray_risk, generate_and_save_gradcam 

app = Flask(__name__)
app.config['UPLOAD_FOLDER'] = 'uploads'
os.makedirs(app.config['UPLOAD_FOLDER'], exist_ok=True)

# ---------------- Load ML Artifacts (Lazy Initialization) ----------------
# Global variables are set to None to force Lazy Loading inside predict()
model = None
target_encoder = None
feature_order = None 

# ---------------- Constants ----------------
CSV_KEYS = [
    'Age', 'Gender', 'Air Pollution', 'Alcohol use', 'Dust Allergy', 'OccuPational Hazards',
    'Genetic Risk', 'chronic Lung Disease', 'Balanced Diet', 'Obesity', 'Smoking',
    'Passive Smoker', 'Chest Pain', 'Coughing of Blood', 'Fatigue', 'Weight Loss',
    'Shortness of Breath', 'Wheezing', 'Swallowing Difficulty', 'Clubbing of Finger Nails',
    'Frequent Cold', 'Dry Cough', 'Snoring'
]

# ---------------- Routes ----------------
@app.route('/')
def home():
    """Render main interactive UI"""
    return render_template('index.html')


@app.route('/predict', methods=['POST'])
def predict():
    """Handles the prediction request and renders the result page."""
    
    # 🛑 CRITICAL FIX: LAZY LOAD MODELS HERE 🛑
    global model, target_encoder, feature_order
    
    if model is None:
        try:
            model = pickle.load(open("model.pkl", "rb"))
            target_encoder = pickle.load(open("target_encoder.pkl", "rb"))
            feature_order = pickle.load(open("model_features.pkl", "rb"))
        except Exception as e:
            return render_template('result.html', final_risk_level="Error", combined_score=0, error_message=f"Initial model loading failed on request: {str(e)}")

    # 1. Check for final Model Load Error
    if not model:
        return render_template('result.html', final_risk_level="Error", combined_score=0, error_message="Model loading failed. Please check server logs.")

    try:
        # --- 1️⃣ Collect and sanitize form data ---
        data_dict = {}
        for key in CSV_KEYS:
            form_key = key.lower().replace(' ', '_')
            val = request.form.get(form_key)
            data_dict[key] = float(val) if val else 0.0

        # --- 2️⃣ Tabular model prediction ---
        features = [data_dict.get(col, 0.0) for col in feature_order]
        X = np.array(features).reshape(1, -1)
        
        if X.shape[1] != len(feature_order):
             raise ValueError("Feature count mismatch. Model features != Form inputs.")
             
        encoded_pred = model.predict(X)[0]
        tabular_proba = model.predict_proba(X)[0]
        tabular_confidence = float(np.max(tabular_proba) * 100)

        tabular_classes = target_encoder.classes_
        high_idx = np.where(tabular_classes == 'High')[0][0] if 'High' in tabular_classes else -1
        tabular_high_prob = float(tabular_proba[high_idx]) if high_idx != -1 else 0.0

        # --- 3️⃣ OCR Risk (PDF / Image) ---
        ocr_risk_score = 0.0
        ocr_keywords = []
        report_file = request.files.get('report')
        if report_file and report_file.filename:
            report_path = os.path.join(app.config['UPLOAD_FOLDER'], report_file.filename)
            report_file.save(report_path)
            extracted_text = extract_keywords_from_report(report_path)
            ocr_risk_score, ocr_keywords = score_text_for_risk(extracted_text)

        # --- 4️⃣ CNN Risk (X-ray) and Grad-CAM Generation ---
        cnn_risk_score = 0.0
        gradcam_url = None 
        xray_file = request.files.get('xray')

        if xray_file and xray_file.filename:
            xray_path = os.path.join(app.config['UPLOAD_FOLDER'], xray_file.filename)
            xray_file.save(xray_path)
            
            abs_xray_path = os.path.abspath(xray_path)

            cnn_risk_score = float(predict_xray_risk(abs_xray_path))
            gradcam_filename = generate_and_save_gradcam(abs_xray_path) 
            
            # 3. REVERTED URL ASSIGNMENT (Using filename directly)
            if gradcam_filename:
                gradcam_url = gradcam_filename # 👈 REVERTED TO FILENAME ASSIGNMENT

        # --- 5️⃣ Risk Fusion Logic ---
        WEIGHT_TABULAR = 0.5
        WEIGHT_OCR = 0.3
        WEIGHT_CNN = 0.2
        combined_score = (
            WEIGHT_TABULAR * tabular_high_prob +
            WEIGHT_OCR * ocr_risk_score +
            WEIGHT_CNN * cnn_risk_score
        )

        if combined_score >= 0.65:
            final_risk = "High"
        elif combined_score >= 0.35:
            final_risk = "Medium"
        else:
            final_risk = "Low"

        # --- 6️⃣ Render Template Response ---
        return render_template('result.html', 
                               final_risk_level=final_risk, 
                               confidence=round(tabular_confidence, 2),
                               user_data=data_dict,
                               ocr_score=round(ocr_risk_score * 100, 2),
                               ocr_keywords=ocr_keywords,
                               cnn_score=round(cnn_risk_score * 100, 2),
                               combined_score=round(combined_score * 100, 2),
                               gradcam_image_url=gradcam_url)

    except Exception as e:
        return render_template('result.html', 
                               final_risk_level="Error", 
                               combined_score=0, 
                               error_message=f"Processing Error: {str(e)}",
                               user_data={})


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