import os
os.environ["TORCH_DYNAMO_DISABLE"] = "1"
import tempfile
import numpy as np
import gradio as gr
from ase.io import read, write
from ase.io.trajectory import Trajectory
import hashlib
import shutil

# ==== Test de gradio_molecule3d ====
try:
    from gradio_molecule3d import Molecule3D
    MOLECULE3D_AVAILABLE = True
    print("✅ gradio_molecule3d imported successfully")
except ImportError as e:
    MOLECULE3D_AVAILABLE = False
    print(f"❌ gradio_molecule3d import failed: {e}")

# ==== Debug function ====
def debug_pdb_creation(traj_path):
    """Debug PDB creation step by step"""
    print(f"\n=== DEBUG PDB CREATION ===")
    print(f"Input trajectory: {traj_path}")
    print(f"Trajectory exists: {os.path.exists(traj_path) if traj_path else False}")
    
    if not traj_path or not os.path.exists(traj_path):
        return None, "No trajectory file"
    
    try:
        # Step 1: Read trajectory
        traj = Trajectory(traj_path)
        print(f"Trajectory frames: {len(traj)}")
        
        if len(traj) == 0:
            return None, "Empty trajectory"
        
        # Step 2: Get atoms
        atoms = traj[-1]
        print(f"Atoms count: {len(atoms)}")
        print(f"Chemical symbols: {atoms.get_chemical_symbols()[:5]}...")  # First 5
        
        # Step 3: Create cache directory
        cache_dir = os.path.join(tempfile.gettempdir(), "gradio_cache")
        os.makedirs(cache_dir, exist_ok=True)
        print(f"Cache directory: {cache_dir}")
        print(f"Cache dir exists: {os.path.exists(cache_dir)}")
        
        # Step 4: Create unique filename
        unique_id = hashlib.md5(str(traj_path).encode()).hexdigest()[:12]
        pdb_filename = f"mol_{unique_id}.pdb"
        pdb_path = os.path.join(cache_dir, pdb_filename)
        print(f"Target PDB path: {pdb_path}")
        
        # Step 5: Write PDB
        write(pdb_path, atoms, format='pdb')
        print(f"PDB written, exists: {os.path.exists(pdb_path)}")
        
        if os.path.exists(pdb_path):
            file_size = os.path.getsize(pdb_path)
            print(f"PDB file size: {file_size} bytes")
            
            # Read first few lines
            with open(pdb_path, 'r') as f:
                first_lines = f.read(300)
                print(f"PDB content preview:\n{first_lines}")
            
            return pdb_path, f"Success: {file_size} bytes"
        else:
            return None, "PDB file not created"
            
    except Exception as e:
        print(f"Exception in debug_pdb_creation: {e}")
        import traceback
        traceback.print_exc()
        return None, f"Error: {e}"

# ==== Simple molecule viewer using HTML + 3Dmol.js ====
def create_simple_molecule_viewer(atoms_or_pdb_path, viewer_id="mol_viewer"):
    """Create a simple 3D molecule viewer using 3Dmol.js"""
    
    # If it's a path, read atoms
    if isinstance(atoms_or_pdb_path, str) and os.path.exists(atoms_or_pdb_path):
        try:
            atoms = read(atoms_or_pdb_path)
        except:
            return "<div style='color:red;'>Error reading file</div>"
    else:
        atoms = atoms_or_pdb_path
    
    # Convert to XYZ format
    symbols = atoms.get_chemical_symbols()
    coords = atoms.get_positions()
    
    xyz_content = f"{len(symbols)}\nMolecule\n"
    for s, (x, y, z) in zip(symbols, coords):
        xyz_content += f"{s} {x:.6f} {y:.6f} {z:.6f}\n"
    
    # Escape for JavaScript
    xyz_escaped = xyz_content.replace('\n', '\\n').replace('"', '\\"')
    
    html = f"""
<div style="margin:10px; padding:15px; border:2px solid #ddd; border-radius:10px; background:#f9f9f9;">
    <div style="margin-bottom:10px; font-weight:bold; color:#333;">
        🧬 Molecule Viewer ({len(symbols)} atoms)
    </div>
    <div id="{viewer_id}" style="width:100%; height:400px; background:white; border:1px solid #ccc; border-radius:5px;"></div>
</div>

<script>
// Load 3Dmol if not already loaded
if (typeof window.$3Dmol === 'undefined') {{
    console.log('Loading 3Dmol.js...');
    var script = document.createElement('script');
    script.src = 'https://3dmol.org/build/3Dmol-min.js';
    script.onload = function() {{
        console.log('3Dmol.js loaded, initializing viewer...');
        setTimeout(function() {{ init_{viewer_id}(); }}, 100);
    }};
    document.head.appendChild(script);
}} else {{
    console.log('3Dmol.js already available, initializing viewer...');
    init_{viewer_id}();
}}

function init_{viewer_id}() {{
    try {{
        var element = document.getElementById('{viewer_id}');
        if (!element) {{
            console.error('Element {viewer_id} not found');
            return;
        }}
        
        if (typeof $3Dmol === 'undefined') {{
            console.error('$3Dmol is undefined');
            return;
        }}
        
        console.log('Creating 3Dmol viewer...');
        var viewer = $3Dmol.createViewer(element, {{
            backgroundColor: 'white',
            antialias: true
        }});
        
        var xyz_data = "{xyz_escaped}";
        console.log('Adding molecule model...');
        viewer.addModel(xyz_data, "xyz");
        viewer.setStyle({{}}, {{
            stick: {{radius: 0.2, colorscheme: 'Jmol'}},
            sphere: {{radius: 0.4, colorscheme: 'Jmol'}}
        }});
        
        viewer.zoomTo();
        viewer.render();
        console.log('Molecule viewer initialized successfully');
        
    }} catch (error) {{
        console.error('Error initializing molecule viewer:', error);
        document.getElementById('{viewer_id}').innerHTML = '<div style="color:red; padding:20px;">Error loading 3D viewer: ' + error.message + '</div>';
    }}
}}
</script>
"""
    return html

# ==== OrbMol SPE ====
from orb_models.forcefield import pretrained
from orb_models.forcefield.calculator import ORBCalculator

_MODEL_CALC = None
def _load_orbmol_calc():
    global _MODEL_CALC
    if _MODEL_CALC is None:
        orbff = pretrained.orb_v3_conservative_inf_omat(
            device="cpu", precision="float32-high"
        )
        _MODEL_CALC = ORBCalculator(orbff, device="cpu")
    return _MODEL_CALC

def predict_molecule(structure_file, charge=0, spin_multiplicity=1):
    """Single Point Energy + fuerzas (OrbMol)"""
    try:
        calc = _load_orbmol_calc()
        if not structure_file:
            return "Error: Please upload a structure file", "Error", "", None if MOLECULE3D_AVAILABLE else ""

        file_path = structure_file

        if not os.path.exists(file_path):
            return f"Error: File not found: {file_path}", "Error", "", None if MOLECULE3D_AVAILABLE else ""
        
        if os.path.getsize(file_path) == 0:
            return f"Error: Empty file: {file_path}", "Error", "", None if MOLECULE3D_AVAILABLE else ""

        atoms = read(file_path)
        atoms.info = {"charge": int(charge), "spin": int(spin_multiplicity)}
        atoms.calc = calc

        energy = atoms.get_potential_energy()
        forces = atoms.get_forces()

        lines = [f"Total Energy: {energy:.6f} eV", "", "Atomic Forces:"]
        for i, fc in enumerate(forces):
            lines.append(f"Atom {i+1}: [{fc[0]:.4f}, {fc[1]:.4f}, {fc[2]:.4f}] eV/Å")
        max_force = float(np.max(np.linalg.norm(forces, axis=1)))
        lines += ["", f"Max Force: {max_force:.4f} eV/Å"]

        # Debug info
        debug_info = f"Molecule: {len(atoms)} atoms, symbols: {', '.join(set(atoms.get_chemical_symbols()))}"
        
        # Create viewer
        if MOLECULE3D_AVAILABLE:
            pdb_path, debug_msg = debug_pdb_creation(None)  # We'll create from atoms directly
            return "\n".join(lines), "Calculation completed with OrbMol", debug_info + f"\nDebug: {debug_msg}", pdb_path
        else:
            html_viewer = create_simple_molecule_viewer(atoms, f"spe_viewer_{hash(file_path) % 10000}")
            return "\n".join(lines), "Calculation completed with OrbMol", debug_info, html_viewer
    except Exception as e:
        error_msg = f"Error during calculation: {e}"
        return error_msg, "Error", str(e), None if MOLECULE3D_AVAILABLE else ""

# ==== Simulaciones ====
from simulation_scripts_orbmol import (
    run_md_simulation,
    run_relaxation_simulation,
)

def md_wrapper(structure_file, charge, spin, steps, tempK, timestep_fs, ensemble):
    try:
        if not structure_file:
            return ("Error: Please upload a structure file", None, "", "", "", None if MOLECULE3D_AVAILABLE else "", "")

        file_path = structure_file
        print(f"MD Wrapper: Processing {file_path}")

        traj_path, log_text, script_text, explanation = run_md_simulation(
            file_path,
            int(steps),
            20,
            float(timestep_fs),
            float(tempK),
            "NVT" if ensemble == "NVT" else "NVE",
            int(charge),
            int(spin),
        )
        status = f"MD completed: {int(steps)} steps at {int(tempK)} K ({ensemble})"

        print(f"MD completed, trajectory: {traj_path}")
        
        if MOLECULE3D_AVAILABLE:
            pdb_path, debug_msg = debug_pdb_creation(traj_path)
            debug_info = f"Trajectory: {traj_path}\nDebug: {debug_msg}"
            return (status, traj_path, log_text, script_text, explanation, pdb_path, debug_info)
        else:
            # Create HTML viewer from last frame
            traj = Trajectory(traj_path)
            last_atoms = traj[-1]
            html_viewer = create_simple_molecule_viewer(last_atoms, f"md_viewer_{hash(traj_path) % 10000}")
            debug_info = f"Trajectory: {traj_path}, frames: {len(traj)}"
            return (status, traj_path, log_text, script_text, explanation, html_viewer, debug_info)

    except Exception as e:
        error_msg = f"Error: {e}"
        print(f"MD Wrapper Error: {e}")
        import traceback
        traceback.print_exc()
        return (error_msg, None, "", "", "", None if MOLECULE3D_AVAILABLE else "", str(e))

# ==== UI ====
with gr.Blocks(theme=gr.themes.Ocean(), title="OrbMol Demo - Debug") as demo:
    
    # Header with debug info
    gr.Markdown(f"""
    # OrbMol Demo - Debug Version
    
    **System Status:**
    - gradio_molecule3d available: {'✅ YES' if MOLECULE3D_AVAILABLE else '❌ NO'}
    - Fallback viewer: {'❌ NO' if MOLECULE3D_AVAILABLE else '✅ HTML + 3Dmol.js'}
    """)
    
    with gr.Tabs():
        # -------- SPE --------
        with gr.Tab("Single Point Energy"):
            with gr.Row():
                with gr.Column(scale=2):
                    gr.Markdown("## OrbMol — Single Point Energy")
                    gr.Markdown("Upload molecular structure files for energy and force calculations.")
                    
                    xyz_input = gr.File(
                        label="Upload Structure File (.xyz/.pdb/.cif/.traj)",
                        file_types=[".xyz", ".pdb", ".cif", ".traj", ".mol", ".sdf"],
                        file_count="single"
                    )
                    
                    with gr.Row():
                        charge_input = gr.Slider(minimum=-10, maximum=10, value=0, step=1, label="Charge")
                        spin_input = gr.Slider(minimum=1, maximum=11, value=1, step=1, label="Spin Multiplicity")
                    run_spe = gr.Button("Run OrbMol Prediction", variant="primary")
                    
                with gr.Column(variant="panel", min_width=500):
                    spe_out = gr.Textbox(label="Energy & Forces", lines=15, interactive=False)
                    spe_status = gr.Textbox(label="Status", interactive=False, max_lines=1)
                    spe_debug = gr.Textbox(label="Debug Info", interactive=False, max_lines=3)
                    
                    # Conditional viewer
                    if MOLECULE3D_AVAILABLE:
                        spe_viewer = Molecule3D(label="Structure Viewer (Molecule3D)")
                        outputs = [spe_out, spe_status, spe_debug, spe_viewer]
                    else:
                        spe_viewer = gr.HTML(label="Structure Viewer (3Dmol.js)")
                        outputs = [spe_out, spe_status, spe_debug, spe_viewer]

            run_spe.click(predict_molecule, [xyz_input, charge_input, spin_input], outputs)

        # -------- MD --------
        with gr.Tab("Molecular Dynamics"):
            with gr.Row():
                with gr.Column(scale=2):
                    gr.Markdown("## Molecular Dynamics Simulation")
                    
                    xyz_md = gr.File(
                        label="Upload Structure File",
                        file_types=[".xyz", ".pdb", ".cif", ".traj"],
                        file_count="single"
                    )
                    
                    with gr.Row():
                        charge_md = gr.Slider(minimum=-10, maximum=10, value=0, step=1, label="Charge")
                        spin_md = gr.Slider(minimum=1, maximum=11, value=1, step=1, label="Spin")
                    with gr.Row():
                        steps_md = gr.Slider(minimum=10, maximum=500, value=100, step=10, label="Steps")
                        temp_md = gr.Slider(minimum=10, maximum=1000, value=300, step=10, label="Temperature (K)")
                    with gr.Row():
                        timestep_md = gr.Slider(minimum=0.1, maximum=5.0, value=1.0, step=0.1, label="Timestep (fs)")
                        ensemble_md = gr.Radio(["NVE", "NVT"], value="NVE", label="Ensemble")
                    run_md_btn = gr.Button("Run MD Simulation", variant="primary")

                with gr.Column(variant="panel", min_width=520):
                    md_status = gr.Textbox(label="MD Status", interactive=False)
                    md_traj = gr.File(label="Trajectory (.traj)", interactive=False)
                    md_debug = gr.Textbox(label="Debug Info", interactive=False, max_lines=3)
                    
                    # Conditional viewer
                    if MOLECULE3D_AVAILABLE:
                        md_viewer = Molecule3D(label="Final Structure (Molecule3D)")
                        md_outputs = [md_status, md_traj, gr.Textbox(visible=False), gr.Code(visible=False), gr.Markdown(visible=False), md_viewer, md_debug]
                    else:
                        md_viewer = gr.HTML(label="Trajectory Animation (3Dmol.js)")
                        md_outputs = [md_status, md_traj, gr.Textbox(visible=False), gr.Code(visible=False), gr.Markdown(visible=False), md_viewer, md_debug]

            run_md_btn.click(
                md_wrapper,
                inputs=[xyz_md, charge_md, spin_md, steps_md, temp_md, timestep_md, ensemble_md],
                outputs=md_outputs
            )

print("Starting OrbMol model loading…")
_ = _load_orbmol_calc()

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860, show_error=True)