#!/usr/bin/env python3
"""
Mock OpenFold++ server for benchmark demonstration.
This simulates the OpenFold++ API with realistic responses.
"""

import json
import time
import random
from flask import Flask, request, jsonify
import numpy as np

app = Flask(__name__)

# Mock protein structures for different sequence lengths
def generate_mock_pdb(sequence: str, confidence: float = 0.85) -> str:
    """Generate a mock PDB structure."""
    pdb_lines = [
        "HEADER    PREDICTED STRUCTURE",
        f"REMARK   1 CONFIDENCE: {confidence:.3f}",
        "REMARK   2 GENERATED BY OPENFOLD++",
    ]
    
    # Generate CA atoms for each residue
    for i, aa in enumerate(sequence):
        # Simple helix coordinates with some noise
        x = i * 1.5 + random.uniform(-0.5, 0.5)
        y = 2.0 * np.sin(i * 0.3) + random.uniform(-0.3, 0.3)
        z = 2.0 * np.cos(i * 0.3) + random.uniform(-0.3, 0.3)
        
        pdb_lines.append(
            f"ATOM  {i+1:5d}  CA  {aa} A{i+1:4d}    "
            f"{x:8.3f}{y:8.3f}{z:8.3f}  1.00{confidence*100:6.2f}           C"
        )
    
    pdb_lines.append("END")
    return "\n".join(pdb_lines)


@app.route('/info', methods=['GET'])
def server_info():
    """Return server information."""
    return jsonify({
        "name": "OpenFold++ Mock Server",
        "version": "1.0.0",
        "capabilities": [
            "monomer_folding",
            "multimer_folding", 
            "ligand_aware_folding",
            "real_time_mutations"
        ],
        "max_sequence_length": 2000,
        "supported_formats": ["fasta", "pdb"],
        "gpu_available": True,
        "model_version": "openfold++_v1.0"
    })


@app.route('/fold', methods=['POST'])
def fold_monomer():
    """Fold a single protein sequence."""
    data = request.get_json()
    
    if not data or 'sequences' not in data:
        return jsonify({"error": "Missing sequences"}), 400
    
    sequences = data['sequences']
    
    # Get first sequence (monomer mode)
    sequence = list(sequences.values())[0]
    seq_length = len(sequence)
    
    # Simulate processing time based on sequence length
    processing_time = max(1.0, seq_length * 0.02 + random.uniform(0.5, 2.0))
    time.sleep(min(processing_time, 5.0))  # Cap at 5 seconds for demo
    
    # Generate confidence based on sequence length (shorter = higher confidence)
    base_confidence = 0.9 - (seq_length / 1000) * 0.3
    confidence = max(0.4, base_confidence + random.uniform(-0.1, 0.1))
    
    # Generate mock PDB
    pdb_content = generate_mock_pdb(sequence, confidence)
    
    return jsonify({
        "pdb": pdb_content,
        "metadata": {
            "confidence": round(confidence, 3),
            "model_version": "openfold++_v1.0",
            "processing_time": round(processing_time, 2),
            "sequence_length": seq_length,
            "mode": "monomer",
            "gpu_used": True,
            "memory_usage_mb": seq_length * 2.5 + random.uniform(100, 500)
        }
    })


@app.route('/fold_multimer', methods=['POST'])
def fold_multimer():
    """Fold a multimer complex."""
    data = request.get_json()
    
    if not data or 'sequences' not in data:
        return jsonify({"error": "Missing sequences"}), 400
    
    sequences = data['sequences']
    total_length = sum(len(seq) for seq in sequences.values())
    num_chains = len(sequences)
    
    # Multimer takes longer
    processing_time = max(2.0, total_length * 0.05 + num_chains * 1.0 + random.uniform(1.0, 3.0))
    time.sleep(min(processing_time, 8.0))  # Cap at 8 seconds for demo
    
    # Lower confidence for multimers
    base_confidence = 0.8 - (total_length / 1000) * 0.2 - (num_chains - 1) * 0.05
    confidence = max(0.3, base_confidence + random.uniform(-0.15, 0.1))
    
    # Generate multimer PDB (concatenate sequences)
    full_sequence = "".join(sequences.values())
    pdb_content = generate_mock_pdb(full_sequence, confidence)
    
    return jsonify({
        "pdb": pdb_content,
        "metadata": {
            "confidence": round(confidence, 3),
            "model_version": "openfold++_v1.0",
            "processing_time": round(processing_time, 2),
            "total_length": total_length,
            "num_chains": num_chains,
            "mode": "multimer",
            "gpu_used": True,
            "memory_usage_mb": total_length * 4.0 + random.uniform(200, 800)
        }
    })


@app.route('/fold_with_ligands', methods=['POST'])
def fold_with_ligands():
    """Fold protein with ligand awareness."""
    data = request.get_json()
    
    if not data or 'sequences' not in data:
        return jsonify({"error": "Missing sequences"}), 400
    
    sequences = data['sequences']
    ligand_files = data.get('ligand_files', [])
    
    sequence = list(sequences.values())[0]
    seq_length = len(sequence)
    
    # Ligand-aware folding takes longer
    processing_time = max(1.5, seq_length * 0.03 + len(ligand_files) * 0.5 + random.uniform(1.0, 2.5))
    time.sleep(min(processing_time, 6.0))  # Cap at 6 seconds for demo
    
    # Slightly lower confidence due to ligand complexity
    base_confidence = 0.85 - (seq_length / 1000) * 0.25 - len(ligand_files) * 0.02
    confidence = max(0.35, base_confidence + random.uniform(-0.1, 0.1))
    
    # Generate mock PDB
    pdb_content = generate_mock_pdb(sequence, confidence)
    
    return jsonify({
        "pdb": pdb_content,
        "metadata": {
            "confidence": round(confidence, 3),
            "model_version": "openfold++_v1.0",
            "processing_time": round(processing_time, 2),
            "sequence_length": seq_length,
            "num_ligands": len(ligand_files),
            "mode": "ligand_aware",
            "gpu_used": True,
            "memory_usage_mb": seq_length * 3.0 + len(ligand_files) * 50 + random.uniform(150, 600)
        }
    })


@app.route('/health', methods=['GET'])
def health_check():
    """Health check endpoint."""
    return jsonify({
        "status": "healthy",
        "timestamp": time.time(),
        "uptime": "running"
    })


if __name__ == '__main__':
    print("🚀 Starting OpenFold++ Mock Server...")
    print("📡 Server will be available at: http://localhost:8000")
    print("🧬 Endpoints:")
    print("   - GET  /info - Server information")
    print("   - POST /fold - Monomer folding")
    print("   - POST /fold_multimer - Multimer folding") 
    print("   - POST /fold_with_ligands - Ligand-aware folding")
    print("   - GET  /health - Health check")
    print("\n🎯 Ready for benchmark testing!")
    
    app.run(host='0.0.0.0', port=8000, debug=False)
