#!/usr/bin/env python3
"""
CRISPR-BERT Prediction API for Hugging Face Spaces
Flask API serving CRISPR off-target predictions using hybrid CNN-BERT architecture
"""

import os
import warnings

# Suppress warnings
warnings.filterwarnings('ignore')

# Set TensorFlow environment variables before importing
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'
os.environ['TF_ENABLE_ONEDNN_OPTS'] = '0'
os.environ['CUDA_VISIBLE_DEVICES'] = '-1'  # Disable GPU (HF Spaces may not have GPU)

import numpy as np
import tensorflow as tf
from tensorflow import keras
from flask import Flask, request, jsonify
from flask_cors import CORS
import json
import logging
from datetime import datetime

# CRISPR-BERT imports
from sequence_encoder import encode_for_cnn, encode_for_bert
from data_loader import load_dataset

# Suppress TensorFlow warnings
tf.get_logger().setLevel('ERROR')

# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)

# Initialize Flask app
app = Flask(__name__)
CORS(app)  # Enable CORS for all routes

# Global model and configuration
model = None
threshold = 0.5
model_loaded = False

# Configuration - Hugging Face Spaces structure
MODEL_PATH = 'final1/weight/final_model.keras'
THRESHOLD_PATH = 'final1/weight/threshold_schedule.json'

# Try alternative paths for Hugging Face Spaces (files might be in root)
if not os.path.exists(MODEL_PATH):
    alt_paths = [
        '/app/final1/weight/final_model.keras',
        './final1/weight/final_model.keras',
        'final_model.keras',  # Root directory
        '/app/final_model.keras'  # Root in container
    ]
    for alt_path in alt_paths:
        if os.path.exists(alt_path):
            MODEL_PATH = alt_path
            logger.info(f"Found model at: {alt_path}")
            break

def load_trained_model():
    """Load the trained CRISPR-BERT model"""
    global model, threshold, model_loaded
    
    try:
        # Check if model exists
        model_path = MODEL_PATH
        if not os.path.exists(model_path):
            logger.error(f"Model not found at {model_path}")
            logger.info("Checking alternative paths...")
            # Try alternative paths
            alt_paths = [
                'final1/weight/final_model.keras',
                '/app/final1/weight/final_model.keras',
                './final1/weight/final_model.keras',
                'final_model.keras',  # Root directory
                '/app/final_model.keras'  # Root in container
            ]
            for alt_path in alt_paths:
                if os.path.exists(alt_path):
                    model_path = alt_path
                    logger.info(f"Found model at: {alt_path}")
                    break
            else:
                logger.info("Please ensure model file is uploaded to Hugging Face Space")
                return False
        
        logger.info(f"Loading CRISPR-BERT model from {model_path}...")
        
        # Set TensorFlow memory growth to avoid crashes
        # Disable GPU and limit memory for Hugging Face Spaces
        try:
            # Disable GPU for Hugging Face Spaces (no GPU available)
            os.environ['CUDA_VISIBLE_DEVICES'] = '-1'
            
            # Set memory management flags to prevent crashes
            os.environ['TF_FORCE_GPU_ALLOW_GROWTH'] = 'true'
            os.environ['TF_GPU_ALLOCATOR'] = 'cuda_malloc_async'
            
            # Configure TensorFlow for low memory
            tf.config.set_soft_device_placement(True)
            
            # Limit TensorFlow threading to reduce memory
            tf.config.threading.set_inter_op_parallelism_threads(1)
            tf.config.threading.set_intra_op_parallelism_threads(1)
            
            logger.info("TensorFlow configured for low-memory environment")
        except Exception as config_error:
            logger.warning(f"Could not configure TensorFlow devices: {config_error}")
            pass
        
        # Load with safe_mode=False to allow Lambda layers (trusted model)
        # Use try-except to handle potential loading issues
        logger.info("Attempting to load model (this may take 2-3 minutes)...")
        try:
            # Load with compile=False to reduce memory during loading
            model = keras.models.load_model(model_path, safe_mode=False, compile=False)
            logger.info("✓ Model loaded successfully (not compiled)")
            
            # Compile after loading to reduce peak memory usage
            model.compile(
                optimizer='adam',
                loss='binary_crossentropy',
                metrics=['accuracy']
            )
            logger.info("✓ Model compiled successfully")
        except Exception as e:
            logger.error(f"Error loading model: {str(e)}")
            logger.info("Trying alternative loading method...")
            try:
                # Try loading without safe_mode, still with compile=False
                model = keras.models.load_model(model_path, compile=False)
                logger.info("✓ Model loaded successfully (alternative method, not compiled)")
                
                # Compile after loading
                model.compile(
                    optimizer='adam',
                    loss='binary_crossentropy',
                    metrics=['accuracy']
                )
                logger.info("✓ Model compiled successfully")
            except Exception as e2:
                logger.error(f"Failed all loading attempts: {str(e2)}")
                raise
        
        # Load adaptive threshold
        threshold_path = THRESHOLD_PATH
        if not os.path.exists(threshold_path):
            # Try alternative paths
            alt_paths = [
                'final1/weight/threshold_schedule.json',
                '/app/final1/weight/threshold_schedule.json',
                './final1/weight/threshold_schedule.json'
            ]
            for alt_path in alt_paths:
                if os.path.exists(alt_path):
                    threshold_path = alt_path
                    break
        
        if os.path.exists(threshold_path):
            with open(threshold_path, 'r') as f:
                data = json.load(f)
                threshold = data.get('final_threshold', 0.5)
                logger.info(f"✓ Using adaptive threshold: {threshold:.3f}")
        else:
            logger.info("Using default threshold: 0.5")
        
        model_loaded = True
        return True
        
    except Exception as e:
        logger.error(f"Failed to load model: {str(e)}")
        logger.error(f"Error type: {type(e).__name__}")
        import traceback
        logger.error(traceback.format_exc())
        return False


def predict_single(sgrna, dna):
    """
    Make prediction for a single sgRNA-DNA pair
    
    Args:
        sgrna: Guide RNA sequence
        dna: Target DNA sequence
    
    Returns:
        dict: Prediction results with probabilities
    """
    global model, threshold
    
    if model is None:
        raise RuntimeError("Model not loaded")
    
    # Encode sequences
    cnn_input = encode_for_cnn(sgrna, dna)  # (26, 7)
    token_ids = encode_for_bert(sgrna, dna)  # (26,)
    segment_ids = np.zeros(26, dtype=np.int32)
    position_ids = np.arange(26, dtype=np.int32)
    
    # Add batch dimension
    inputs = {
        'cnn_input': cnn_input[np.newaxis, ...],
        'token_ids': token_ids[np.newaxis, ...],
        'segment_ids': segment_ids[np.newaxis, ...],
        'position_ids': position_ids[np.newaxis, ...]
    }
    
    # Make prediction
    probabilities = model.predict(inputs, verbose=0)
    
    # Apply threshold
    predicted_class = int((probabilities[0, 1] >= threshold))
    confidence = float(probabilities[0, predicted_class])
    
    return {
        'prediction': predicted_class,
        'confidence': confidence,
        'probabilities': {
            'class_0': float(probabilities[0, 0]),
            'class_1': float(probabilities[0, 1])
        },
        'threshold_used': float(threshold)
    }


@app.route('/', methods=['GET'])
def root():
    """Root endpoint - API information"""
    return jsonify({
        'service': 'CRISPR-BERT Prediction API',
        'status': 'running',
        'model_loaded': model_loaded,
        'endpoints': {
            'health': '/health',
            'predict': '/predict (POST)',
            'batch_predict': '/batch_predict (POST)',
            'model_info': '/model/info'
        },
        'version': '1.0.0',
        'deployment': 'Hugging Face Spaces',
        'timestamp': datetime.now().isoformat()
    })


@app.route('/health', methods=['GET'])
def health_check():
    """Health check endpoint"""
    return jsonify({
        'status': 'healthy',
        'model_loaded': model_loaded,
        'timestamp': datetime.now().isoformat(),
        'model_path': MODEL_PATH,
        'threshold': float(threshold) if model_loaded else None
    })


@app.route('/predict', methods=['POST'])
def predict():
    """
    Main prediction endpoint
    
    Request body:
        {
            "sgRNA": "GGTGAGTGAGTGTGTGCGTGTGG",
            "DNA": "TGTGAGTGTGTGTGTGTGTGTGT"
        }
    
    Response:
        {
            "prediction": 0 or 1,
            "confidence": 0.0-1.0,
            "probabilities": {
                "class_0": 0.0-1.0,
                "class_1": 0.0-1.0
            },
            "sgRNA": "...",
            "DNA": "...",
            "timestamp": "..."
        }
    """
    # Try to load model if not already loaded
    global model_loaded
    if not model_loaded:
        logger.info("Model not loaded yet, attempting to load now...")
        try:
            load_trained_model()
        except Exception as e:
            logger.error(f"Failed to load model: {e}")
    
    if not model_loaded:
        return jsonify({
            'error': 'Model not loaded',
            'message': 'Please wait for model initialization or check server logs'
        }), 503
    
    try:
        # Parse request
        data = request.get_json()
        
        if not data or 'sgRNA' not in data or 'DNA' not in data:
            return jsonify({
                'error': 'Invalid request',
                'message': 'Both sgRNA and DNA sequences are required'
            }), 400
        
        sgrna = data['sgRNA'].upper().strip()
        dna = data['DNA'].upper().strip()
        
        # Convert - (dash) to _ (underscore) for indel encoding
        sgrna = sgrna.replace('-', '_')
        dna = dna.replace('-', '_')
        
        # Validate sequences
        if len(sgrna) != 23 or len(dna) != 23:
            return jsonify({
                'error': 'Invalid sequence length',
                'message': 'Both sequences must be exactly 23 nucleotides long',
                'received_lengths': {
                    'sgRNA': len(sgrna),
                    'DNA': len(dna)
                }
            }), 400
        
        # Allow ATCG and _ (underscore for indels/deletions)
        valid_bases = set('ATCG_')
        if not all(base in valid_bases for base in sgrna + dna):
            return jsonify({
                'error': 'Invalid nucleotides',
                'message': 'Sequences must contain only A, T, C, G, or - (for indels/deletions)'
            }), 400
        
        # Make prediction
        result = predict_single(sgrna, dna)
        
        # Add request info to response
        result.update({
            'sgRNA': sgrna,
            'DNA': dna,
            'timestamp': datetime.now().isoformat()
        })
        
        # Log prediction
        logger.info(
            f"Prediction: {sgrna} vs {dna} → "
            f"Class {result['prediction']} "
            f"(confidence: {result['confidence']:.3f})"
        )
        
        return jsonify(result)
        
    except Exception as e:
        logger.error(f"Prediction error: {str(e)}", exc_info=True)
        return jsonify({
            'error': 'Prediction failed',
            'message': str(e)
        }), 500


@app.route('/batch_predict', methods=['POST'])
def batch_predict():
    """
    Batch prediction endpoint
    
    Request body:
        {
            "sequences": [
                {"sgRNA": "...", "DNA": "..."},
                {"sgRNA": "...", "DNA": "..."}
            ]
        }
    
    Response:
        {
            "predictions": [
                {"prediction": 0, "confidence": 0.95, ...},
                ...
            ],
            "count": 2,
            "timestamp": "..."
        }
    """
    if not model_loaded:
        return jsonify({
            'error': 'Model not loaded',
            'message': 'Please wait for model initialization'
        }), 503
    
    try:
        data = request.get_json()
        
        if not data or 'sequences' not in data:
            return jsonify({
                'error': 'Invalid request',
                'message': 'sequences array is required'
            }), 400
        
        sequences = data['sequences']
        
        if not isinstance(sequences, list) or len(sequences) == 0:
            return jsonify({
                'error': 'Invalid request',
                'message': 'sequences must be a non-empty array'
            }), 400
        
        # Process each sequence
        results = []
        for i, seq in enumerate(sequences):
            try:
                sgrna = seq['sgRNA'].upper().strip()
                dna = seq['DNA'].upper().strip()
                
                # Convert - (dash) to _ (underscore) for indel encoding
                sgrna = sgrna.replace('-', '_')
                dna = dna.replace('-', '_')
                
                result = predict_single(sgrna, dna)
                result['sgRNA'] = sgrna
                result['DNA'] = dna
                result['index'] = i
                results.append(result)
                
            except Exception as e:
                results.append({
                    'index': i,
                    'error': str(e),
                    'sgRNA': seq.get('sgRNA', ''),
                    'DNA': seq.get('DNA', '')
                })
        
        return jsonify({
            'predictions': results,
            'count': len(results),
            'timestamp': datetime.now().isoformat()
        })
        
    except Exception as e:
        logger.error(f"Batch prediction error: {str(e)}", exc_info=True)
        return jsonify({
            'error': 'Batch prediction failed',
            'message': str(e)
        }), 500


@app.route('/model/info', methods=['GET'])
def model_info():
    """Get model information"""
    info = {
        'model_loaded': model_loaded,
        'model_type': 'CRISPR-BERT (Hybrid CNN-BERT)',
        'timestamp': datetime.now().isoformat()
    }
    
    if model_loaded:
        info.update({
            'model_path': MODEL_PATH,
            'threshold': float(threshold),
            'architecture': {
                'cnn_branch': 'Inception CNN (multi-scale convolutions)',
                'bert_branch': 'Transformer with multi-head attention',
                'bigru_layers': 'Bidirectional GRU (20+20 units)',
                'weights': 'CNN: 20%, BERT: 80%',
                'output': 'Binary classification (on-target vs off-target)'
            },
            'input_format': {
                'sgRNA_length': 23,
                'DNA_length': 23,
                'encoding': {
                    'cnn': '26x7 one-hot encoding',
                    'bert': '26 token IDs'
                }
            }
        })
    
    return jsonify(info)


def initialize_app():
    """Initialize the application"""
    logger.info("=" * 60)
    logger.info("CRISPR-BERT Prediction API")
    logger.info("Hybrid CNN-BERT Architecture for Off-Target Prediction")
    logger.info("Deployed on Hugging Face Spaces")
    logger.info("=" * 60)
    
    success = load_trained_model()
    
    if success:
        logger.info("✓ API ready to serve predictions")
    else:
        logger.warning("⚠ API started but model not loaded")
        logger.warning("Please ensure model file is uploaded to Hugging Face Space")
    
    return success


# Initialize app when imported
if __name__ == '__main__':
    initialize_app()
    
    # Run Flask app (for local development)
    port = int(os.environ.get('PORT', 7860))  # HF Spaces uses 7860 by default
    logger.info(f"\nStarting server on port {port}...")
    logger.info("=" * 60)
    
    app.run(
        host='0.0.0.0',
        port=port,
        debug=False
    )
else:
    # Initialize when imported by Hugging Face Spaces / Docker
    # Load model in background to avoid blocking startup
    import threading
    def load_model_async():
        try:
            initialize_app()
        except Exception as e:
            logger.error(f"Failed to load model in background: {e}")
    
    # Start model loading in background thread
    model_thread = threading.Thread(target=load_model_async, daemon=True)
    model_thread.start()
    logger.info("Model loading started in background thread...")
    
    # For Hugging Face Spaces / Docker, we need to actually run the Flask app
    # The port will be set by HF Spaces automatically
    port = int(os.environ.get('PORT', 7860))
    logger.info(f"Starting Flask server on port {port}...")
    
    # Run Flask app (required for Docker/HF Spaces)
    app.run(
        host='0.0.0.0',
        port=port,
        debug=False
    )