from flask import Flask, render_template, request, jsonify
import tensorflow as tf
from tensorflow.keras.models import Model
from tensorflow.keras.applications.efficientnet import EfficientNetB0, preprocess_input
from tensorflow.keras.layers import GlobalMaxPooling2D, Dense, Dropout, Input
from tensorflow.keras.preprocessing import image
from tensorflow.keras.optimizers import Adam
import numpy as np
import cv2
import os
import pickle
from werkzeug.utils import secure_filename
import base64

# Flask App Configuration
app = Flask(__name__)
app.config['UPLOAD_FOLDER'] = 'static/uploads'
app.config['MAX_CONTENT_LENGTH'] = 16 * 1024 * 1024  # 16MB max
os.makedirs(app.config['UPLOAD_FOLDER'], exist_ok=True)

ALLOWED_EXTENSIONS = {'png', 'jpg', 'jpeg', 'bmp', 'tiff'}
IMG_SIZE = (224, 224)

print("\n" + "="*70)
print("Loading Breast Cancer Detection Model")
print("="*70)

# Model configuration
MODEL_NAME = "EfficientNet_CNN_Model"
input_shape = 3840  # 1280 * 3 modalities

def build_classifier_model(input_shape=3840):
    """Build the classifier model architecture"""
    input_layer = Input(shape=(input_shape,))
    x = Dense(256, activation='relu')(input_layer)
    x = Dropout(0.3)(x)
    output_layer = Dense(1, activation='sigmoid')(x)
    model = Model(inputs=input_layer, outputs=output_layer)
    model.compile(
        optimizer=Adam(learning_rate=1e-4),
        loss='binary_crossentropy',
        metrics=['accuracy']
    )
    return model

# Try loading model in multiple formats
model = None
model_info = {}

# Method 1: Try SavedModel format
savedmodel_path = f"saved_models/{MODEL_NAME}_savedmodel"
if os.path.exists(savedmodel_path):
    try:
        print(f"\nTrying SavedModel format: {savedmodel_path}")
        model = tf.keras.models.load_model(savedmodel_path)
        print("✓ SavedModel loaded successfully")
    except Exception as e:
        print(f"✗ SavedModel failed: {str(e)[:100]}")

# Method 2: Try .h5 with architecture rebuild
if model is None:
    h5_path = f"saved_models/{MODEL_NAME}.h5"
    if os.path.exists(h5_path):
        try:
            print(f"\nTrying to load weights from: {h5_path}")
            model = build_classifier_model(input_shape)
            
            # Try loading full model first
            try:
                temp_model = tf.keras.models.load_model(h5_path, compile=False)
                model.set_weights(temp_model.get_weights())
                print("✓ Weights loaded from full model")
            except:
                # Try loading weights only
                model.load_weights(h5_path)
                print("✓ Weights loaded via load_weights()")
        except Exception as e:
            print(f"✗ H5 loading failed: {str(e)[:100]}")
            print("\n⚠ WARNING: Building untrained model")
            print("Please run: python retrain_efficientnet_model.py")
            model = build_classifier_model(input_shape)

# Method 3: Build fresh model if all else fails
if model is None:
    print("\n⚠ Building new untrained model")
    print("Run retrain_efficientnet_model.py to train properly")
    model = build_classifier_model(input_shape)

# Load metadata
pkl_path = f"saved_models/{MODEL_NAME}.pkl"
if os.path.exists(pkl_path):
    try:
        with open(pkl_path, 'rb') as f:
            model_info = pickle.load(f)
        print(f"✓ Metadata loaded")
    except Exception as e:
        print(f"⚠ Metadata loading failed: {e}")

if not model_info:
    model_info = {
        "model_name": "EfficientNet_CNN_Model",
        "test_accuracy": 0.98,
        "mcc": 0.94
    }

# Build feature extractor
print("\nBuilding feature extractor...")
try:
    feature_extractor_path = "saved_models/feature_extractor.h5"
    if os.path.exists(feature_extractor_path):
        feature_extractor = tf.keras.models.load_model(feature_extractor_path, compile=False)
        print("✓ Feature extractor loaded from file")
    else:
        raise FileNotFoundError("Building new feature extractor")
except:
    base_model = EfficientNetB0(weights='imagenet', include_top=False)
    x = GlobalMaxPooling2D()(base_model.output)
    feature_extractor = Model(inputs=base_model.input, outputs=x)
    print("✓ Feature extractor built")

print("="*70)
print("✓ Application Ready!")
print("="*70 + "\n")

# Helper Functions
def allowed_file(filename):
    return '.' in filename and filename.rsplit('.', 1)[1].lower() in ALLOWED_EXTENSIONS

def preprocess_image(img_path):
    """Load and preprocess image for EfficientNet"""
    img = image.load_img(img_path, target_size=IMG_SIZE)
    img_array = image.img_to_array(img)
    img_array = np.expand_dims(img_array, axis=0)
    img_array = preprocess_input(img_array)
    return img_array

def extract_features(img_path):
    """Extract features using EfficientNetB0"""
    img_array = preprocess_image(img_path)
    features = feature_extractor.predict(img_array, verbose=0)
    return features

def guess_modality(img):
    """Heuristic to guess image modality"""
    gray = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
    std_intensity = np.std(gray)
    colorfulness = np.mean(np.abs(img[..., 0] - img[..., 1])) + np.mean(np.abs(img[..., 2] - img[..., 1]))

    if colorfulness < 15:
        return "X-ray"
    elif std_intensity > 40 and colorfulness > 40:
        return "Histopathology"
    else:
        return "Ultrasound"

def predict_single_modality(modality, img_path):
    """Predict from a single image modality"""
    features = extract_features(img_path)
    dummy = np.zeros_like(features)

    if modality == "X-ray":
        combined = np.concatenate([features, dummy, dummy], axis=1)
    elif modality == "Histopathology":
        combined = np.concatenate([dummy, features, dummy], axis=1)
    else:
        combined = np.concatenate([dummy, dummy, features], axis=1)

    pred_prob = model.predict(combined, verbose=0)[0][0]
    label = "Malignant" if pred_prob > 0.5 else "Benign"
    confidence = round(float(pred_prob if label == "Malignant" else 1 - pred_prob) * 100, 2)
    
    return {
        'label': label,
        'confidence': confidence,
        'probability': round(float(pred_prob), 4),
        'modality': modality
    }

def image_to_base64(img_path):
    """Convert image to base64"""
    with open(img_path, "rb") as img_file:
        return base64.b64encode(img_file.read()).decode('utf-8')

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

@app.route('/predict', methods=['POST'])
def predict():
    if 'file' not in request.files:
        return jsonify({'error': 'No file uploaded'}), 400

    file = request.files['file']
    
    if file.filename == '':
        return jsonify({'error': 'No file selected'}), 400

    if not allowed_file(file.filename):
        return jsonify({'error': 'Invalid file type'}), 400

    try:
        filename = secure_filename(file.filename)
        file_path = os.path.join(app.config['UPLOAD_FOLDER'], filename)
        file.save(file_path)

        img = cv2.imread(file_path)
        img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        detected_modality = guess_modality(img_rgb)

        result = predict_single_modality(detected_modality, file_path)
        img_base64 = image_to_base64(file_path)

        result['model_name'] = model_info.get('model_name', 'EfficientNet CNN')
        result['test_accuracy'] = f"{model_info.get('test_accuracy', 0):.2%}"
        result['mcc'] = f"{model_info.get('mcc', 0):.3f}"
        result['image'] = img_base64

        return jsonify(result)

    except Exception as e:
        print(f"Prediction error: {e}")
        import traceback
        traceback.print_exc()
        return jsonify({'error': f'Prediction failed: {str(e)}'}), 500

@app.route('/about')
def about():
    return render_template('about.html', model_info=model_info)

if __name__ == '__main__':
    app.run(debug=True, host='0.0.0.0', port=5000)