import os
import torch
import torch.nn as nn
import torch.nn.functional as F
from torchvision import models, transforms
from PIL import Image
import numpy as np
import cv2
from flask import Flask, request, jsonify
from flask_cors import CORS
from flasgger import Swagger, swag_from
from ultralytics import YOLO
import io
import base64
import logging

# Load environment variables from .env file if it exists
try:
    from dotenv import load_dotenv
    load_dotenv()
except ImportError:
    pass

# Production-ready logging configuration
DEBUG_MODE = os.getenv('DEBUG', 'False').lower() == 'true'
log_level = logging.DEBUG if DEBUG_MODE else logging.INFO
logging.basicConfig(level=log_level, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

app = Flask(__name__)
CORS(app)

# Swagger configuration
swagger_config = {
    "headers": [],
    "specs": [
        {
            "endpoint": 'apispec',
            "route": '/apispec.json',
            "rule_filter": lambda rule: True,
            "model_filter": lambda tag: True,
        }
    ],
    "static_url_path": "/flasgger_static",
    "swagger_ui": True,
    "specs_route": "/docs"
}

swagger_template = {
    "swagger": "2.0",
    "info": {
        "title": "Lung Cancer Classification API with Grad-CAM",
        "description": "API for classifying lung cancer types using DenseNet121 with Grad-CAM visualization",
        "version": "1.0.0",
        "contact": {
            "name": "API Support"
        }
    },
    "host": "localhost:5001",
    "basePath": "/",
    "schemes": ["http"],
    "consumes": ["multipart/form-data", "application/json"],
    "produces": ["application/json", "image/jpeg"]
}

swagger = Swagger(app, config=swagger_config, template=swagger_template)

# --- 1. CONFIGURATION ---
MODEL_PATH = 'models/densenet_final_classification.pth'
YOLO_MODEL_PATH = 'models/best.pt'
NUM_CLASSES = 4
PADDING_FACTOR = 0.20  # 20% context margin around detected tumors
LABELS = [
    'Adenocarcinoma (Class A)',
    'Small Cell (Class B)',
    'Large Cell (Class E)',
    'Squamous Cell (Class G)'
]

# --- 2. MODEL SETUP ---
class GradCAM:
    def __init__(self, model, target_layer):
        self.model = model
        self.target_layer = target_layer
        self.activations = None
        
        self.target_layer.register_forward_hook(self._save_activations)

    def _save_activations(self, module, input, output):
        self.activations = output

    def generate_heatmap(self, input_image, class_idx=None):
        self.model.eval()
        self.activations = None

        output = self.model(input_image)

        if class_idx is None:
            class_idx = torch.argmax(output).item()

        # Compute gradients directly w.r.t. activations — avoids backward hook + view issue
        grads = torch.autograd.grad(
            outputs=output[0, class_idx],
            inputs=self.activations,
            retain_graph=False,
            create_graph=False
        )[0]

        # Pool gradients over spatial dimensions
        pooled_gradients = torch.mean(grads, dim=[0, 2, 3])

        # Weight activations by pooled gradients
        activations = self.activations.detach()
        weighted = torch.sum(
            activations * pooled_gradients.view(1, -1, 1, 1),
            dim=1
        ).squeeze()

        # Apply ReLU and normalize
        heatmap = F.relu(weighted)
        max_val = torch.max(heatmap)
        if max_val > 0:
            heatmap = heatmap / max_val

        return heatmap.cpu().numpy(), LABELS[class_idx], torch.softmax(output, dim=1)[0][class_idx].item()

def load_model():
    # DenseNet121 is the common backbone for these tasks
    logger.info("🧠 Loading DenseNet121 Model...")
    model = models.densenet121(weights=None)
    num_ftrs = model.classifier.in_features  # 1024
    
    # Classifier matching densenet_final_classification.pth:
    # State dict has parameters at index 1 (Linear layer)
    # Index 0 is a non-parameter layer (likely ReLU as per README)
    model.classifier = nn.Sequential(
        nn.ReLU(),                              # classifier.0
        nn.Linear(num_ftrs, NUM_CLASSES)        # classifier.1 (has weight and bias)
    )
    
    if os.path.exists(MODEL_PATH):
        model.load_state_dict(torch.load(MODEL_PATH, map_location=torch.device('cpu')))
        logger.info(f"✅ Model loaded from {MODEL_PATH}")
        
        # === NEW: Verify weight distribution ===
        classifier = model.classifier
        
        # Check output layer (final Linear layer at index 1)
        output_layer = classifier[1]  # nn.Linear(num_ftrs, 4)
        final_bias = output_layer.bias.data
        final_weight = output_layer.weight.data
        
        logger.warning("\n" + "="*60)
        logger.warning("🔍 MODEL WEIGHT ANALYSIS (At Startup)")
        logger.warning("="*60)
        
        logger.info(f"Output layer bias values: {[f'{x.item():.4f}' for x in final_bias]}")
        logger.info("Output layer weight stats:")
        for i, label in enumerate(LABELS):
            weight_mean = final_weight[i].mean().item()
            weight_std = final_weight[i].std().item()
            bias_val = final_bias[i].item()
            logger.info(f"  {label}: weight_mean={weight_mean:.4f}, weight_std={weight_std:.4f}, bias={bias_val:.4f}")
        
        # Check for extreme imbalance
        class_b_bias = final_bias[1].item()  # Class B is index 1
        class_b_weight_mean = final_weight[1].mean().item()
        
        if class_b_bias > 1.5 or class_b_weight_mean > 0.3:
            logger.warning("⚠️  CLASS B (SMALL CELL) BIAS DETECTED!")
            logger.warning(f"    Bias value: {class_b_bias:.4f} (should be close to other classes)")
            logger.warning(f"    Weight mean: {class_b_weight_mean:.4f}")
            logger.warning("    This explains why all images are classified as Class B.")
            logger.warning("    Root cause: Model trained on imbalanced data or needs retraining.")
        
        logger.warning("="*60 + "\n")
        
    else:
        logger.warning(f"⚠️ Model file not found at {MODEL_PATH}. Using uninitialized model.")
    
    model.eval()
    return model

model = load_model()
# Load YOLO model for tumor detection
logger.info("👁️ Loading YOLO Detection Model...")
if os.path.exists(YOLO_MODEL_PATH):
    yolo_model = YOLO(YOLO_MODEL_PATH)
    logger.info(f"✅ YOLO model loaded from {YOLO_MODEL_PATH}")
else:
    logger.warning(f"⚠️ YOLO model file not found at {YOLO_MODEL_PATH}")
    yolo_model = None

# DenseNet target layer is usually the last feature block
target_layer = model.features.norm5 
cam = GradCAM(model, target_layer)

# Confidence threshold — below this we treat the result as uncertain
CONFIDENCE_THRESHOLD = 0.5

# --- 3. IMAGE PREPROCESSING ---
# Normalize with ImageNet stats (standard for DenseNet pretrained backbone)
preprocess = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])

preprocess_no_norm = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
])

def is_ct_scan(image_pil):
    """
    Validate that the image is likely a CT scan.
    CT scans are grayscale — R, G, B channels are nearly identical.
    Color photos (faces, etc.) have high inter-channel variance.
    """
    img_np = np.array(image_pil.resize((64, 64))).astype(np.float32)
    
    r, g, b = img_np[:,:,0], img_np[:,:,1], img_np[:,:,2]
    
    # Mean absolute difference between channels
    rg_diff = np.mean(np.abs(r - g))
    rb_diff = np.mean(np.abs(r - b))
    gb_diff = np.mean(np.abs(g - b))
    color_score = (rg_diff + rb_diff + gb_diff) / 3.0
    
    # CT scans are grayscale: channel diff < threshold
    # Real CT scans: typically 2-5
    # Color photos (faces): 7-100
    # Threshold: 6.0
    is_valid = color_score < 6.0
    
    if is_valid:
        logger.info(f"✅ Valid CT scan detected (color_score={color_score:.2f})")
    else:
        logger.warning(f"❌ Not a CT scan (color_score={color_score:.2f} >= 6.0)")
    
    return is_valid, round(float(color_score), 2)


def apply_heatmap(image_pil, heatmap):
    # Resize heatmap to match image size
    heatmap = cv2.resize(heatmap, (image_pil.size[0], image_pil.size[1]))
    heatmap = np.uint8(255 * heatmap)
    heatmap = cv2.applyColorMap(heatmap, cv2.COLORMAP_JET)
    
    img_np = np.array(image_pil)
    # Convert RGB (PyTorch/PIL) to BGR (OpenCV)
    img_np = cv2.cvtColor(img_np, cv2.COLOR_RGB2BGR)
    
    superimposed_img = heatmap * 0.4 + img_np
    superimposed_img = np.clip(superimposed_img, 0, 255).astype(np.uint8)
    
    # Convert BGR (OpenCV) back to RGB (PyTorch/PIL)
    return cv2.cvtColor(superimposed_img, cv2.COLOR_BGR2RGB)

def apply_heatmap_to_full_image(image_pil, heatmap_crop, crop_coords):
    """
    Apply heatmap from cropped tumor region to full image.
    
    Args:
        image_pil: Full original image (PIL)
        heatmap_crop: Heatmap generated from crop (numpy array)
        crop_coords: (crop_x1, crop_y1, crop_x2, crop_y2) coordinates
    
    Returns:
        Full image with heatmap overlay
    """
    crop_x1, crop_y1, crop_x2, crop_y2 = crop_coords
    img_h, img_w = image_pil.size[1], image_pil.size[0]
    
    # Create full-size heatmap initialized to zeros
    full_heatmap = np.zeros((img_h, img_w), dtype=np.float32)
    
    # Resize cropped heatmap to match crop size
    crop_h = crop_y2 - crop_y1
    crop_w = crop_x2 - crop_x1
    resized_heatmap = cv2.resize(heatmap_crop, (crop_w, crop_h))
    
    # Place resized heatmap at correct location in full image
    full_heatmap[crop_y1:crop_y2, crop_x1:crop_x2] = resized_heatmap
    
    # Normalize to 0-1 range
    max_val = np.max(full_heatmap)
    if max_val > 0:
        full_heatmap = full_heatmap / max_val
    
    # Apply color mapping
    full_heatmap = np.uint8(255 * full_heatmap)
    colored_heatmap = cv2.applyColorMap(full_heatmap, cv2.COLORMAP_JET)
    
    img_np = np.array(image_pil)
    # Convert RGB (PyTorch/PIL) to BGR (OpenCV)
    img_np = cv2.cvtColor(img_np, cv2.COLOR_RGB2BGR)
    
    # Blend with original image (40% heatmap, 60% original)
    superimposed_img = colored_heatmap.astype(np.float32) * 0.4 + img_np.astype(np.float32) * 0.6
    superimposed_img = np.clip(superimposed_img, 0, 255).astype(np.uint8)
    
    # Convert BGR (OpenCV) back to RGB (PyTorch/PIL)
    return cv2.cvtColor(superimposed_img, cv2.COLOR_BGR2RGB)

# --- 4. ROUTES ---
@app.route('/', methods=['GET'])
def home():
    """
    Home endpoint
    ---
    tags:
      - General
    responses:
      200:
        description: API information and available endpoints
        schema:
          type: object
          properties:
            status:
              type: string
              example: running
            message:
              type: string
              example: Lung Cancer Classification API with Grad-CAM
            endpoints:
              type: object
            model:
              type: string
              example: models/densenet_final_classification.pth
            classes:
              type: array
              items:
                type: string
    """
    return jsonify({
        'status': 'running',
        'message': 'Lung Cancer Classification API with Grad-CAM',
        'endpoints': {
            '/health': 'GET - Check API health and model status',
            '/validate-ct': 'POST - Check if image is a CT scan (no classification)',
            '/analyze': 'POST - Upload image for Grad-CAM analysis (multipart/form-data with "file" field)',
            '/docs': 'GET - Swagger UI documentation'
        },
        'model': MODEL_PATH,
        'classes': LABELS
    })

@app.route('/validate-ct', methods=['POST'])
def validate_ct():
    """
    Validate if uploaded image is a CT scan WITHOUT running classification
    ---
    tags:
      - Validation
    parameters:
      - name: file
        in: formData
        type: file
        required: true
        description: Medical image file (JPEG, PNG, etc.)
    consumes:
      - multipart/form-data
    produces:
      - application/json
    responses:
      200:
        description: CT scan validation result
        schema:
          type: object
          properties:
            is_ct_scan:
              type: boolean
              example: true
            color_score:
              type: number
              format: float
              example: 5.2
            message:
              type: string
              example: Valid CT scan - grayscale image detected
    """
    logger.info("=== /validate-ct endpoint called ===")
    
    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
    
    try:
        logger.info(f"Validating file: {file.filename}")
        img_bytes = file.read()
        image = Image.open(io.BytesIO(img_bytes)).convert('RGB')
        logger.info(f"Image loaded, size: {image.size}, mode: {image.mode}")
        
        valid_ct, color_score = is_ct_scan(image)
        
        # Convert numpy types to Python native types for JSON serialization
        valid_ct = bool(valid_ct)
        color_score = float(color_score)
        
        logger.info(f"Validation result: {valid_ct}, color_score: {color_score}")
        
        return jsonify({
            'is_ct_scan': valid_ct,
            'color_score': color_score,
            'message': f"{'Valid CT scan - grayscale image detected' if valid_ct else f'NOT a CT scan - color image detected (color_score={color_score})'}"
        })
    
    except Exception as e:
        logger.error(f"Error validating image: {str(e)}", exc_info=True)
        return jsonify({'error': str(e)}), 500

@app.route('/analyze', methods=['POST'])
def analyze():
    """
    Analyze lung cancer image with YOLO detection + DenseNet classification + Grad-CAM
    ---
    tags:
      - Analysis
    parameters:
      - name: file
        in: formData
        type: file
        required: true
        description: Medical image file (JPEG, PNG, etc.)
    consumes:
      - multipart/form-data
    produces:
      - application/json
    responses:
      200:
        description: Successfully analyzed image
        schema:
          type: object
          properties:
            success:
              type: boolean
              example: true
            tumors_detected:
              type: integer
              example: 1
            detections:
              type: array
              items:
                type: object
                properties:
                  tumor_id:
                    type: integer
                  bbox:
                    type: array
                    items: [x1, y1, x2, y2]
                  prediction:
                    type: string
                  confidence:
                    type: number
                  all_confidences:
                    type: object
                  crop_image:
                    type: string
                  heatmap_image:
                    type: string
            original_image:
              type: string
              description: Base64 encoded original image with detection boxes
      400:
        description: Bad request or no tumors detected
    """
    logger.warning("\n" + "="*80)
    logger.warning("=== /analyze endpoint called (YOLO + DenseNet + Grad-CAM) ===")
    logger.warning("="*80)
    
    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
    
    try:
        logger.info(f"📁 Processing file: {file.filename}")
        img_bytes = file.read()
        
        # Load image as BGR (OpenCV format)
        original_img = cv2.imdecode(np.frombuffer(img_bytes, np.uint8), cv2.IMREAD_COLOR)
        if original_img is None:
            return jsonify({'error': 'Could not load image'}), 400
        
        h, w = original_img.shape[:2]
        logger.info(f"✅ Image loaded - Size: {w}x{h}")

        # STEP 1: YOLO DETECTION
        logger.info("\n>>> STEP 1: YOLO Tumor Detection")
        if yolo_model is None:
            return jsonify({'error': 'YOLO model not loaded'}), 500
        
        results = yolo_model.predict(source=original_img, device='cpu', conf=0.15, verbose=False)
        boxes = results[0].boxes
        
        logger.info(f"🔍 Detected {len(boxes)} tumor(s)")
        
        if len(boxes) == 0:
            logger.warning("ℹ️ No tumors detected")
            # Still return the original image
            pil_img = Image.fromarray(cv2.cvtColor(original_img, cv2.COLOR_BGR2RGB))
            original_io = io.BytesIO()
            pil_img.save(original_io, 'JPEG', quality=85)
            original_io.seek(0)
            original_base64 = base64.b64encode(original_io.getvalue()).decode('utf-8')
            
            return jsonify({
                'success': True,
                'tumors_detected': 0,
                'detections': [],
                'original_image': original_base64,
                'message': 'No tumors detected in this image'
            })

        # STEP 2: CLASSIFY EACH DETECTED TUMOR
        logger.info("\n>>> STEP 2: Classifying Detected Tumors")
        detections = []
        segmentation_img = original_img.copy()
        
        # Track highest confidence tumor for heatmap display
        highest_conf_idx = -1
        highest_conf_value = 0
        highest_conf_heatmap = None
        highest_conf_crop_pil = None
        highest_conf_crop_coords = None
        
        for tumor_idx, box in enumerate(boxes):
            logger.info(f"\n--- Processing Tumor {tumor_idx + 1} ---")
            
            # Convert numpy types to Python ints for JSON serialization
            x1, y1, x2, y2 = [int(v) for v in box.xyxy[0].numpy()]
            
            # Add padding
            box_w, box_h = x2 - x1, y2 - y1
            pad_w = int(box_w * PADDING_FACTOR)
            pad_h = int(box_h * PADDING_FACTOR)
            crop_x1 = int(max(0, x1 - pad_w))
            crop_y1 = int(max(0, y1 - pad_h))
            crop_x2 = int(min(w, x2 + pad_w))
            crop_y2 = int(min(h, y2 + pad_h))
            
            logger.info(f"  Box: [{x1}, {y1}, {x2}, {y2}]")
            logger.info(f"  Crop (with padding): [{crop_x1}, {crop_y1}, {crop_x2}, {crop_y2}]")
            
            # Crop tumor region from color image
            tumor_crop_bgr = original_img[crop_y1:crop_y2, crop_x1:crop_x2]
            
            # Convert BGR (OpenCV) to RGB (PyTorch)
            tumor_crop_rgb = cv2.cvtColor(tumor_crop_bgr, cv2.COLOR_BGR2RGB)
            pil_crop = Image.fromarray(tumor_crop_rgb)
            
            # Try both preprocessing approaches
            logger.info("  Running inference with NORMALIZED preprocessing...")
            input_norm = preprocess(pil_crop).unsqueeze(0)
            
            logger.info("  Running inference with NON-NORMALIZED preprocessing...")
            input_nonorm = preprocess_no_norm(pil_crop).unsqueeze(0)
            
            with torch.no_grad():
                logits_norm = model(input_norm)[0]
                logits_nonorm = model(input_nonorm)[0]
                out_norm = torch.softmax(logits_norm, dim=0)
                out_nonorm = torch.softmax(logits_nonorm, dim=0)
            
            max_conf_norm = torch.max(out_norm).item()
            max_conf_nonorm = torch.max(out_nonorm).item()
            
            logger.info(f"  Normalized max confidence: {max_conf_norm*100:.2f}%")
            logger.info(f"  Non-normalized max confidence: {max_conf_nonorm*100:.2f}%")
            
            # DEBUG: Log raw logits
            logger.debug(f"  Raw logits (normalized): {[f'{x.item():.2f}' for x in logits_norm]}")
            logger.debug(f"  Raw logits (non-normalized): {[f'{x.item():.2f}' for x in logits_nonorm]}")
            
            # Use whichever gives higher confidence
            if max_conf_norm >= max_conf_nonorm:
                input_tensor = input_norm
                probs = out_norm
                logits_used = logits_norm
                logger.info("  ✅ Using NORMALIZED preprocessing")
            else:
                input_tensor = input_nonorm
                probs = out_nonorm
                logits_used = logits_nonorm
                logger.info("  ✅ Using NON-NORMALIZED preprocessing")
            
            confidence = torch.max(probs).item()
            class_idx = torch.argmax(probs).item()
            diagnosis = LABELS[class_idx]
            
            logger.warning(f"  🎯 Diagnosis: {diagnosis}")
            logger.warning(f"  📈 Confidence: {confidence*100:.2f}%")
            
            # All class confidences
            all_confidences = {
                LABELS[i]: round(probs[i].item() * 100, 2)
                for i in range(NUM_CLASSES)
            }
            
            logger.info("  📋 Full class breakdown:")
            for i, (label, conf) in enumerate(all_confidences.items()):
                logit_val = logits_used[i].item()
                logger.info(f"      - {label}: {conf}% (logit: {logit_val:.3f})")
            
            # ⚠️ BIAS DETECTION: Warn if Class B (index 1) is always winning
            if class_idx == 1:  # Class B is index 1
                logger.warning(f"  ⚠️  WARNING: Class B detected (index 1). Check for model bias.")
            
            # Generate Grad-CAM heatmap for this tumor
            logger.info("  Generating Grad-CAM heatmap...")
            heatmap_np, _, _ = cam.generate_heatmap(input_tensor, class_idx=class_idx)
            
            # Track highest confidence tumor for main heatmap display
            if confidence > highest_conf_value:
                highest_conf_value = confidence
                highest_conf_idx = tumor_idx
                highest_conf_heatmap = heatmap_np
                highest_conf_crop_pil = pil_crop
                highest_conf_crop_coords = (crop_x1, crop_y1, crop_x2, crop_y2)
            
            # Apply heatmap to crop
            result_img = apply_heatmap(pil_crop, heatmap_np)
            result_pil = Image.fromarray(result_img)
            
            # Encode crop and heatmap as base64
            crop_io = io.BytesIO()
            pil_crop.save(crop_io, 'JPEG', quality=85)
            crop_io.seek(0)
            crop_base64 = base64.b64encode(crop_io.getvalue()).decode('utf-8')
            
            heatmap_io = io.BytesIO()
            result_pil.save(heatmap_io, 'JPEG', quality=85)
            heatmap_io.seek(0)
            heatmap_base64 = base64.b64encode(heatmap_io.getvalue()).decode('utf-8')
            
            # Draw box on segmentation image for visualization
            cv2.rectangle(segmentation_img, (x1, y1), (x2, y2), (0, 255, 0), 2)
            label = f"Tumor {tumor_idx + 1}: {diagnosis.split(' ')[0]} ({confidence*100:.1f}%)"
            cv2.putText(segmentation_img, label, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 255, 0), 2)
            
            detections.append({
                'tumor_id': tumor_idx + 1,
                'bbox': [int(x1), int(y1), int(x2), int(y2)],
                'bbox_with_padding': [crop_x1, crop_y1, crop_x2, crop_y2],
                'prediction': diagnosis,
                'confidence': round(confidence * 100, 2),
                'all_confidences': all_confidences,
                'crop_image': crop_base64,
                'heatmap_image': heatmap_base64
            })
        
        # Encode original image with detection boxes
        detection_pil = Image.fromarray(cv2.cvtColor(segmentation_img, cv2.COLOR_BGR2RGB))
        detection_io = io.BytesIO()
        detection_pil.save(detection_io, 'JPEG', quality=85)
        detection_io.seek(0)
        detection_base64 = base64.b64encode(detection_io.getvalue()).decode('utf-8')
        
        # Encode heatmap of highest confidence tumor on FULL IMAGE
        full_original_pil = Image.fromarray(cv2.cvtColor(original_img, cv2.COLOR_BGR2RGB))
        heatmap_display_img = apply_heatmap_to_full_image(
            full_original_pil, 
            highest_conf_heatmap, 
            highest_conf_crop_coords
        )
        heatmap_display_pil = Image.fromarray(heatmap_display_img)
        heatmap_display_io = io.BytesIO()
        heatmap_display_pil.save(heatmap_display_io, 'JPEG', quality=85)
        heatmap_display_io.seek(0)
        heatmap_display_base64 = base64.b64encode(heatmap_display_io.getvalue()).decode('utf-8')
        
        logger.warning("\n" + "="*80)
        logger.warning("✅ ANALYSIS COMPLETE - RETURNING SUCCESS RESPONSE")
        logger.warning("="*80)
        
        # Classification summary
        class_counts = {}
        for detection in detections:
            pred = detection['prediction'].split(' ')[0]  # Get first word (class name)
            class_counts[pred] = class_counts.get(pred, 0) + 1
        
        logger.warning("📊 CLASSIFICATION SUMMARY:")
        for class_name, count in sorted(class_counts.items()):
            pct = (count / len(detections)) * 100 if detections else 0
            logger.warning(f"   {class_name}: {count} tumor(s) ({pct:.1f}%)")
        
        if len(detections) > 0 and 'Small' in class_counts:
            pct_class_b = (class_counts.get('Small', 0) / len(detections)) * 100
            if pct_class_b >= 80:
                logger.warning(f"⚠️  HIGH CLASS B BIAS DETECTED: {pct_class_b:.0f}% classified as Small Cell!")
        
        logger.warning("="*80 + "\n")
        
        return jsonify({
            'success': True,
            'tumors_detected': len(boxes),
            'detections': detections,
            'detection_image': detection_base64,
            'heatmap_image': heatmap_display_base64
        })

    except Exception as e:
        logger.error(f"\n❌ ERROR processing image: {str(e)}", exc_info=True)
        return jsonify({'error': str(e)}), 500

@app.route('/health', methods=['GET'])
def health():
    """
    Health check endpoint
    ---
    tags:
      - General
    responses:
      200:
        description: API health status
        schema:
          type: object
          properties:
            status:
              type: string
              example: healthy
            model_loaded:
              type: boolean
              example: true
    """
    return jsonify({'status': 'healthy', 'model_loaded': os.path.exists(MODEL_PATH)})

if __name__ == '__main__':
    port = int(os.getenv('PORT', 5001))
    app.run(host='0.0.0.0', port=port, debug=DEBUG_MODE)