app.py

"""
DNA-Diffusion Gradio Application with Integrated 3D Viewer
Combines Gradio interface with HTML-based 3D molecular visualization
"""

import gradio as gr
import logging
import json
import os
from typing import Dict, Any, Tuple, List
import time
import random
import numpy as np
from datetime import datetime
import urllib.parse

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

# Import spaces and handle ZeroGPU
try:
    import spaces
    SPACES_AVAILABLE = True
    logger.info("Spaces module loaded successfully")
except ImportError:
    SPACES_AVAILABLE = False
    logger.warning("Spaces module not available - running without GPU")
    # Create dummy decorator
    class spaces:
        @staticmethod
        def GPU(duration=60):
            def decorator(func):
                return func
            return decorator

# DNA Model and genetic code table (same as before)
CODON_TABLE = {
    'TTT': 'F', 'TTC': 'F', 'TTA': 'L', 'TTG': 'L',
    'TCT': 'S', 'TCC': 'S', 'TCA': 'S', 'TCG': 'S',
    'TAT': 'Y', 'TAC': 'Y', 'TAA': '*', 'TAG': '*',
    'TGT': 'C', 'TGC': 'C', 'TGA': '*', 'TGG': 'W',
    'CTT': 'L', 'CTC': 'L', 'CTA': 'L', 'CTG': 'L',
    'CCT': 'P', 'CCC': 'P', 'CCA': 'P', 'CCG': 'P',
    'CAT': 'H', 'CAC': 'H', 'CAA': 'Q', 'CAG': 'Q',
    'CGT': 'R', 'CGC': 'R', 'CGA': 'R', 'CGG': 'R',
    'ATT': 'I', 'ATC': 'I', 'ATA': 'I', 'ATG': 'M',
    'ACT': 'T', 'ACC': 'T', 'ACA': 'T', 'ACG': 'T',
    'AAT': 'N', 'AAC': 'N', 'AAA': 'K', 'AAG': 'K',
    'AGT': 'S', 'AGC': 'S', 'AGA': 'R', 'AGG': 'R',
    'GTT': 'V', 'GTC': 'V', 'GTA': 'V', 'GTG': 'V',
    'GCT': 'A', 'GCC': 'A', 'GCA': 'A', 'GCG': 'A',
    'GAT': 'D', 'GAC': 'D', 'GAA': 'E', 'GAG': 'E',
    'GGT': 'G', 'GGC': 'G', 'GGA': 'G', 'GGG': 'G'
}

class DNAModel:
    """Mock DNA generation model"""
    def generate(self, cell_type: str, num_sequences: int = 1, length: int = 200, guidance_scale: float = 1.0):
        sequences = []
        
        for _ in range(num_sequences):
            sequence = ''.join(random.choice(['A', 'T', 'C', 'G']) for _ in range(length))
            
            if cell_type == "K562":
                for i in range(0, length-8, 50):
                    sequence = sequence[:i] + 'GCGCGCGC' + sequence[i+8:]
            elif cell_type == "GM12878":
                for i in range(10, length-8, 60):
                    sequence = sequence[:i] + 'ATATATAT' + sequence[i+8:]
            elif cell_type == "HepG2":
                for i in range(20, length-12, 70):
                    sequence = sequence[:i] + 'GCGATCGATCGC' + sequence[i+12:]
            
            sequences.append(sequence)
        
        return sequences[0] if num_sequences == 1 else sequences

model = DNAModel()

class DNADiffusionApp:
    def __init__(self):
        self.current_sequence = ""
        self.current_analysis = {}
        self.generation_count = 0
        
    @spaces.GPU(duration=60)
    def generate_with_gpu(self, cell_type: str, guidance_scale: float = 1.0, use_enhanced: bool = True):
        logger.info(f"Generating sequence on GPU for cell type: {cell_type}")
        
        try:
            time.sleep(2)
            sequence = model.generate(cell_type, length=200, guidance_scale=guidance_scale)
            
            if use_enhanced:
                sequence = self.enhance_sequence(sequence, cell_type)
            
            self.generation_count += 1
            logger.info(f"Successfully generated sequence #{self.generation_count}")
            
            return sequence
            
        except Exception as e:
            logger.error(f"GPU generation failed: {e}")
            raise
    
    def enhance_sequence(self, sequence: str, cell_type: str) -> str:
        enhancers = {
            "K562": "GGGACTTTCC",
            "GM12878": "TGACGTCA",
            "HepG2": "TGTTGGTGG"
        }
        
        if cell_type in enhancers:
            pos = len(sequence) // 4
            enhancer = enhancers[cell_type]
            sequence = sequence[:pos] + enhancer + sequence[pos+len(enhancer):]
        
        return sequence
    
    def analyze_sequence(self, sequence: str) -> Dict[str, Any]:
        if not sequence:
            return {}
        
        gc_count = sequence.count('G') + sequence.count('C')
        gc_content = (gc_count / len(sequence)) * 100
        
        if len(sequence) < 14:
            tm = 4 * (sequence.count('G') + sequence.count('C')) + 2 * (sequence.count('A') + sequence.count('T'))
        else:
            tm = 81.5 + 0.41 * gc_content - 675 / len(sequence)
        
        restriction_sites = {}
        enzymes = {
            'EcoRI': 'GAATTC',
            'BamHI': 'GGATCC',
            'HindIII': 'AAGCTT',
            'PstI': 'CTGCAG',
            'XbaI': 'TCTAGA'
        }
        
        for enzyme, pattern in enzymes.items():
            positions = []
            for i in range(len(sequence) - len(pattern) + 1):
                if sequence[i:i+len(pattern)] == pattern:
                    positions.append(i)
            if positions:
                restriction_sites[enzyme] = positions
        
        protein = self.translate_to_protein(sequence)
        
        return {
            'length': len(sequence),
            'gc_content': round(gc_content, 1),
            'melting_temp': round(tm, 1),
            'restriction_sites': restriction_sites,
            'protein_length': len(protein),
            'protein_sequence': protein[:50] + '...' if len(protein) > 50 else protein
        }
    
    def translate_to_protein(self, dna_sequence: str) -> str:
        protein = []
        
        start_pos = dna_sequence.find('ATG')
        if start_pos == -1:
            start_pos = 0
        
        for i in range(start_pos, len(dna_sequence) - 2, 3):
            codon = dna_sequence[i:i+3]
            if len(codon) == 3:
                amino_acid = CODON_TABLE.get(codon, 'X')
                if amino_acid == '*':
                    break
                protein.append(amino_acid)
        
        return ''.join(protein)

app = DNADiffusionApp()

def load_html_file(filename):
    """Safely load and prepare HTML file for iframe embedding"""
    try:
        with open(filename, 'r', encoding='utf-8') as f:
            content = f.read()
        # Properly encode content for data URI
        encoded_content = urllib.parse.quote(content, safe='')
        return f'<iframe src="data:text/html;charset=utf-8,{encoded_content}" width="100%" height="800px" style="border: none; border-radius: 10px;"></iframe>'
    except FileNotFoundError:
        return f'<div style="padding: 40px; text-align: center; color: #ff6b6b; background: #2a2a2a; border-radius: 10px;"><h3>⚠️ File Not Found</h3><p>{filename} not found in the current directory.</p></div>'
    except Exception as e:
        return f'<div style="padding: 40px; text-align: center; color: #ff6b6b; background: #2a2a2a; border-radius: 10px;"><h3>❌ Error Loading File</h3><p>{str(e)}</p></div>'

# HTML for 3D Viewer
HTML_3D_VIEWER = """
<!DOCTYPE html>
<html>
<head>
    <meta charset="UTF-8">
    <style>
        body {
            margin: 0;
            padding: 0;
            background: #000;
            font-family: Arial, sans-serif;
            color: #fff;
            height: 100vh;
            overflow: hidden;
        }
        #viewer-container {
            width: 100%;
            height: 100%;
            position: relative;
        }
        #canvas3d {
            width: 100%;
            height: 100%;
        }
        .controls-panel {
            position: absolute;
            top: 20px;
            right: 20px;
            background: rgba(0,0,0,0.8);
            padding: 20px;
            border-radius: 10px;
            border: 2px solid #00ff88;
            max-width: 300px;
        }
        .controls-panel h3 {
            color: #00ff88;
            margin-top: 0;
        }
        .control-btn {
            background: #00ff88;
            color: #000;
            border: none;
            padding: 8px 16px;
            margin: 5px;
            border-radius: 5px;
            cursor: pointer;
            font-weight: bold;
        }
        .control-btn:hover {
            background: #00cc66;
        }
        .info-display {
            position: absolute;
            bottom: 20px;
            left: 20px;
            background: rgba(0,0,0,0.8);
            padding: 15px;
            border-radius: 8px;
            border: 1px solid #0088ff;
        }
        #sequence-display {
            font-family: monospace;
            color: #00ff88;
            word-break: break-all;
            margin-top: 10px;
        }
    </style>
</head>
<body>
    <div id="viewer-container">
        <canvas id="canvas3d"></canvas>
        
        <div class="controls-panel">
            <h3>3D View Controls</h3>
            <button class="control-btn" onclick="setViewMode('cartoon')">Cartoon</button>
            <button class="control-btn" onclick="setViewMode('stick')">Stick</button>
            <button class="control-btn" onclick="setViewMode('sphere')">Sphere</button>
            <button class="control-btn" onclick="toggleRotation()">Toggle Rotation</button>
            <button class="control-btn" onclick="resetView()">Reset View</button>
        </div>
        
        <div class="info-display">
            <strong>Current Sequence:</strong>
            <div id="sequence-display">No sequence loaded</div>
        </div>
    </div>

    <script src="https://cdnjs.cloudflare.com/ajax/libs/three.js/r128/three.min.js"></script>
    <script>
        let scene, camera, renderer;
        let moleculeGroup;
        let currentSequence = '';
        let autoRotate = true;
        let viewMode = 'cartoon';

        function init() {
            scene = new THREE.Scene();
            scene.background = new THREE.Color(0x000000);
            
            camera = new THREE.PerspectiveCamera(75, window.innerWidth / window.innerHeight, 0.1, 1000);
            camera.position.z = 50;
            
            renderer = new THREE.WebGLRenderer({ canvas: document.getElementById('canvas3d'), antialias: true });
            renderer.setSize(window.innerWidth, window.innerHeight);
            
            const ambientLight = new THREE.AmbientLight(0x404040, 1.5);
            scene.add(ambientLight);
            
            const directionalLight = new THREE.DirectionalLight(0xffffff, 1);
            directionalLight.position.set(50, 50, 50);
            scene.add(directionalLight);
            
            moleculeGroup = new THREE.Group();
            scene.add(moleculeGroup);
            
            animate();
            
            // Listen for sequence updates from parent
            window.addEventListener('message', function(e) {
                if (e.data.type === 'updateSequence') {
                    updateSequence(e.data.sequence);
                }
            });
        }

        function updateSequence(sequence) {
            currentSequence = sequence;
            document.getElementById('sequence-display').textContent = sequence;
            generateDNAStructure(sequence);
        }

        function generateDNAStructure(sequence) {
            moleculeGroup.clear();
            
            const radius = 10;
            const rise = 3.4;
            const basesPerTurn = 10;
            const anglePerBase = (2 * Math.PI) / basesPerTurn;

            // Create double helix
            const curve1Points = [];
            const curve2Points = [];
            
            for (let i = 0; i < sequence.length; i++) {
                const angle = i * anglePerBase;
                const height = i * rise / basesPerTurn;
                
                curve1Points.push(new THREE.Vector3(
                    radius * Math.cos(angle),
                    height,
                    radius * Math.sin(angle)
                ));
                
                curve2Points.push(new THREE.Vector3(
                    radius * Math.cos(angle + Math.PI),
                    height,
                    radius * Math.sin(angle + Math.PI)
                ));
            }
            
            // Create backbone curves
            const curve1 = new THREE.CatmullRomCurve3(curve1Points);
            const curve2 = new THREE.CatmullRomCurve3(curve2Points);
            
            const tubeGeometry1 = new THREE.TubeGeometry(curve1, 100, 0.5, 8, false);
            const tubeGeometry2 = new THREE.TubeGeometry(curve2, 100, 0.5, 8, false);
            
            const material1 = new THREE.MeshPhongMaterial({ color: 0xff0000 });
            const material2 = new THREE.MeshPhongMaterial({ color: 0x0000ff });
            
            moleculeGroup.add(new THREE.Mesh(tubeGeometry1, material1));
            moleculeGroup.add(new THREE.Mesh(tubeGeometry2, material2));
            
            // Add base pairs
            for (let i = 0; i < Math.min(sequence.length, 50); i++) {
                const p1 = curve1Points[i];
                const p2 = curve2Points[i];
                
                const direction = new THREE.Vector3().subVectors(p2, p1);
                const distance = direction.length();
                direction.normalize();
                
                const geometry = new THREE.CylinderGeometry(0.3, 0.3, distance, 8);
                const material = new THREE.MeshPhongMaterial({ 
                    color: getBaseColor(sequence[i])
                });
                
                const cylinder = new THREE.Mesh(geometry, material);
                cylinder.position.copy(p1).add(direction.multiplyScalar(distance / 2));
                cylinder.quaternion.setFromUnitVectors(new THREE.Vector3(0, 1, 0), direction);
                
                moleculeGroup.add(cylinder);
            }
            
            // Center the molecule
            const box = new THREE.Box3().setFromObject(moleculeGroup);
            const center = box.getCenter(new THREE.Vector3());
            moleculeGroup.position.sub(center);
        }

        function getBaseColor(base) {
            const colors = {
                'A': 0xff0000,
                'T': 0x00ff00,
                'G': 0x0000ff,
                'C': 0xffff00
            };
            return colors[base] || 0xffffff;
        }

        function setViewMode(mode) {
            viewMode = mode;
            if (currentSequence) {
                generateDNAStructure(currentSequence);
            }
        }

        function toggleRotation() {
            autoRotate = !autoRotate;
        }

        function resetView() {
            camera.position.set(0, 0, 50);
            moleculeGroup.rotation.set(0, 0, 0);
        }

        function animate() {
            requestAnimationFrame(animate);
            
            if (autoRotate) {
                moleculeGroup.rotation.y += 0.01;
            }
            
            renderer.render(scene, camera);
        }

        // Mouse controls
        let isDragging = false;
        let previousMousePosition = { x: 0, y: 0 };

        document.addEventListener('mousedown', function(e) {
            isDragging = true;
        });

        document.addEventListener('mouseup', function(e) {
            isDragging = false;
        });

        document.addEventListener('mousemove', function(e) {
            if (isDragging) {
                const deltaMove = {
                    x: e.clientX - previousMousePosition.x,
                    y: e.clientY - previousMousePosition.y
                };

                moleculeGroup.rotation.y += deltaMove.x * 0.01;
                moleculeGroup.rotation.x += deltaMove.y * 0.01;
            }

            previousMousePosition = {
                x: e.clientX,
                y: e.clientY
            };
        });

        document.addEventListener('wheel', function(e) {
            camera.position.z += e.deltaY * 0.1;
            camera.position.z = Math.max(10, Math.min(200, camera.position.z));
        });

        window.addEventListener('resize', function() {
            camera.aspect = window.innerWidth / window.innerHeight;
            camera.updateProjectionMatrix();
            renderer.setSize(window.innerWidth, window.innerHeight);
        });

        // Initialize
        init();
    </script>
</body>
</html>
"""

def create_demo():
    """Create the Gradio interface with integrated 3D viewer"""
    
    css = """
    .gradio-container {
        font-family: 'Arial', sans-serif;
        background: linear-gradient(135deg, #0a0a0a 0%, #1a1a1a 100%);
    }
    .sequence-box {
        font-family: 'Courier New', monospace;
        background-color: #1e1e1e;
        color: #00ff88;
        padding: 15px;
        border-radius: 8px;
        border: 1px solid #00ff88;
    }
    iframe {
        border: none;
        border-radius: 10px;
        box-shadow: 0 4px 6px rgba(0, 0, 0, 0.1);
    }
    .gr-button-primary {
        background: linear-gradient(135deg, #00ff88, #0088ff) !important;
        border: none !important;
        color: #000 !important;
    }
    .gr-button-secondary {
        background: linear-gradient(135deg, #ff6b6b, #ff9f40) !important;
        border: none !important;
        color: #fff !important;
    }
    """
    
    with gr.Blocks(css=css, title="DNA-Diffusion Suite", theme=gr.themes.Soft()) as demo:
        gr.Markdown(
            """
            # 🧬 DNA-Diffusion Suite: AI-Powered Molecular Engineering
            
            Generate, visualize, and gamify DNA sequences with advanced AI models!
            
            [![GPU Status](https://img.shields.io/badge/GPU-Enabled-brightgreen)](https://huggingface.co/spaces)
            """
        )
        
        gpu_status = gr.Markdown(
            f"🖥️ **GPU Status**: {'✅ Available' if SPACES_AVAILABLE else '❌ Not Available (CPU mode)'}"
        )
        
        with gr.Tabs():
            # Tab 1: Sequence Generation
            with gr.TabItem("🎯 Generate Sequence"):
                with gr.Row():
                    with gr.Column(scale=1):
                        cell_type = gr.Dropdown(
                            choices=["K562", "GM12878", "HepG2"],
                            value="K562",
                            label="Cell Type",
                            info="Select the cell type for sequence generation"
                        )
                        
                        guidance_scale = gr.Slider(
                            minimum=1.0,
                            maximum=10.0,
                            value=1.0,
                            step=0.5,
                            label="Guidance Scale",
                            info="Higher values create more cell-type specific patterns"
                        )
                        
                        enhanced_mode = gr.Checkbox(
                            value=True,
                            label="Enhanced Mode",
                            info="Add cell-type specific enhancer sequences"
                        )
                        
                        generate_btn = gr.Button(
                            "🚀 Generate Sequence",
                            variant="primary",
                            size="lg"
                        )
                    
                    with gr.Column(scale=2):
                        sequence_output = gr.Textbox(
                            label="Generated DNA Sequence",
                            lines=5,
                            max_lines=10,
                            elem_classes="sequence-box"
                        )
                        
                        analysis_output = gr.JSON(
                            label="Sequence Analysis"
                        )
            
            # Tab 2: 3D Visualization
            with gr.TabItem("🔬 3D Structure"):
                with gr.Row():
                    with gr.Column():
                        gr.Markdown("### 3D DNA Structure Visualization")
                        gr.Markdown("The 3D viewer shows the double helix structure of your generated DNA sequence.")
                        
                        # HTML component for 3D viewer
                        encoded_viewer = urllib.parse.quote(HTML_3D_VIEWER, safe='')
                        viewer_html = gr.HTML(
                            value=f'<iframe src="data:text/html;charset=utf-8,{encoded_viewer}" width="100%" height="600px"></iframe>',
                            label="3D Molecular Viewer"
                        )
                        
                        # Button to update 3D view
                        update_3d_btn = gr.Button(
                            "🔄 Update 3D View with Current Sequence",
                            variant="secondary"
                        )
            
            # Tab 3: DNA Casino (if external HTML exists)
            with gr.TabItem("🎰 DNA Casino"):
                gr.Markdown("### DNA Slot Machine Game")
                gr.Markdown("Experience DNA generation as a casino game!")
                
                # Load external HTML file using helper function
                casino_html = gr.HTML(
                    value=load_html_file('dna-slot-machine.html'),
                    label="DNA Casino Game"
                )
                
                # Button to capture sequence from casino
                with gr.Row():
                    casino_sequence = gr.Textbox(
                        label="Casino Generated Sequence",
                        placeholder="Sequence from casino will appear here",
                        interactive=False
                    )
                    analyze_casino_btn = gr.Button("🔬 Analyze Casino Sequence", variant="secondary")
            
            # Tab 4: Batch Generation
            with gr.TabItem("📦 Batch Generation"):
                with gr.Row():
                    with gr.Column():
                        batch_cell_types = gr.CheckboxGroup(
                            choices=["K562", "GM12878", "HepG2"],
                            value=["K562", "GM12878"],
                            label="Cell Types for Batch"
                        )
                        
                        batch_count = gr.Slider(
                            minimum=1,
                            maximum=10,
                            value=5,
                            step=1,
                            label="Number of Sequences"
                        )
                        
                        batch_generate_btn = gr.Button(
                            "🚀 Generate Batch",
                            variant="primary"
                        )
                    
                    with gr.Column():
                        batch_output = gr.Dataframe(
                            headers=["ID", "Cell Type", "Length", "GC%", "Tm(°C)"],
                            label="Batch Results"
                        )
        
        status_text = gr.Textbox(
            label="Status",
            value="Ready to generate sequences...",
            interactive=False
        )
        
        # Event handlers
        def generate_single(cell_type, guidance_scale, enhanced):
            try:
                status_text.value = "🔄 Generating sequence on GPU..."
                
                sequence = app.generate_with_gpu(cell_type, guidance_scale, enhanced)
                analysis = app.analyze_sequence(sequence)
                
                status_text.value = f"✅ Successfully generated sequence for {cell_type}"
                return sequence, analysis, status_text.value
                
            except Exception as e:
                error_msg = f"❌ Error: {str(e)}"
                logger.error(error_msg)
                return "", {}, error_msg
        
        def update_3d_viewer(sequence):
            if not sequence:
                return viewer_html.value
            
            # Create HTML with embedded sequence data
            html_with_sequence = HTML_3D_VIEWER.replace(
                "window.addEventListener('message'",
                f"updateSequence('{sequence}');\n        window.addEventListener('message'"
            )
            
            # Properly encode content for data URI
            encoded_content = urllib.parse.quote(html_with_sequence, safe='')
            return f'<iframe src="data:text/html;charset=utf-8,{encoded_content}" width="100%" height="600px"></iframe>'
        
        def generate_batch(cell_types, count):
            if not cell_types:
                return None, "❌ Please select at least one cell type"
            
            try:
                status_text.value = "🔄 Generating batch on GPU..."
                
                results = []
                for i in range(count):
                    cell_type = cell_types[i % len(cell_types)]
                    sequence = app.generate_with_gpu(cell_type)
                    analysis = app.analyze_sequence(sequence)
                    
                    results.append([
                        f'SEQ_{i+1:03d}',
                        cell_type,
                        analysis['length'],
                        analysis['gc_content'],
                        analysis['melting_temp']
                    ])
                
                status_text.value = f"✅ Generated {len(results)} sequences"
                return results, status_text.value
                
            except Exception as e:
                error_msg = f"❌ Error: {str(e)}"
                logger.error(error_msg)
                return None, error_msg
        
        # Connect event handlers
        generate_btn.click(
            fn=generate_single,
            inputs=[cell_type, guidance_scale, enhanced_mode],
            outputs=[sequence_output, analysis_output, status_text]
        )
        
        update_3d_btn.click(
            fn=update_3d_viewer,
            inputs=[sequence_output],
            outputs=[viewer_html]
        )
        
        # Auto-update 3D viewer when sequence is generated
        sequence_output.change(
            fn=update_3d_viewer,
            inputs=[sequence_output],
            outputs=[viewer_html]
        )
        
        batch_generate_btn.click(
            fn=generate_batch,
            inputs=[batch_cell_types, batch_count],
            outputs=[batch_output, status_text]
        )
        
        # Casino sequence analysis
        def analyze_casino_sequence(seq):
            if not seq:
                return {}, "❌ No sequence to analyze"
            
            analysis = app.analyze_sequence(seq)
            return analysis, "✅ Casino sequence analyzed"
        
        try:
            analyze_casino_btn.click(
                fn=analyze_casino_sequence,
                inputs=[casino_sequence],
                outputs=[analysis_output, status_text]
            )
        except NameError:
            # Casino tab components may not be defined if HTML file is missing
            pass
    
    return demo

if __name__ == "__main__":
    logger.info("=" * 50)
    logger.info("DNA-Diffusion App with 3D Viewer Starting...")
    logger.info(f"GPU Available: {SPACES_AVAILABLE}")
    logger.info(f"Environment: {'Hugging Face Spaces' if os.getenv('SPACE_ID') else 'Local'}")
    logger.info("=" * 50)
    
    demo = create_demo()
    
    if os.getenv("SPACE_ID"):
        demo.launch(
            server_name="0.0.0.0",
            server_port=7860,
            share=False,
            show_error=True
        )
    else:
        demo.launch(
            share=True,
            show_error=True,
            debug=True
        )