src/Home.py

import streamlit as st
import os
import tempfile
import numpy as np
import py3Dmol
import streamlit.components.v1 as components
from io import StringIO
import time
from pyscf import gto, dft, scf
from pyscf.dft import gen_grid
import re
import base64
import contextlib
import sys
import io as _io
from ase import Atoms
from ase.io import read as ase_read
import pandas as pd

# Set page configuration
st.set_page_config(
    page_title='PyFock GUI - Interactive DFT Calculations',
    layout='wide',
    page_icon="⚛️",
    menu_items={
        'About': "PyFock GUI - A web interface for PyFock, a pure Python DFT code with Numba JIT acceleration"
    }
)

# Sidebar with enhanced styling
st.sidebar.image("https://raw.githubusercontent.com/manassharma07/PyFock/main/logo_crysx_pyfock.png", use_container_width=True)

st.sidebar.markdown("---")

# About PyFock section
st.sidebar.markdown("### About PyFock")
st.sidebar.markdown("""
**Pure Python DFT** with performance matching C++ codes!

**Key Advantages:**
- 100% Pure Python (including molecular integrals)
- Numba JIT acceleration
- GPU support (CUDA via CuPy)
- Near-quadratic scaling (~O(N²·⁰⁵))
- Accuracy matching PySCF (<10⁻⁷ Ha)
- Windows/Linux/MacOS compatible
- Easy pip installation
""")

st.sidebar.markdown("---")

# Features
st.sidebar.markdown("### GUI Features")
st.sidebar.markdown("""
* Run DFT in your browser
* Visualize HOMO, LUMO, density
* Compare with PySCF
* Download cube files & scripts
* Interactive 3D visualization
* No installation required!
""")

st.sidebar.markdown("---")

# Links section
st.sidebar.markdown("### 🔗 Links & Resources")
st.sidebar.markdown("""
[![GitHub](https://img.shields.io/badge/GitHub-Repository-blue?logo=github)](https://github.com/manassharma07/PyFock)
[![PyPI](https://img.shields.io/badge/PyPI-Package-orange?logo=pypi)](https://pypi.org/project/pyfock/)
[![Docs](https://img.shields.io/badge/Documentation-Read-green?logo=readthedocs)](https://github.com/manassharma07/PyFock/blob/main/Documentation.md)

📄 **Paper:** [arXiv preprint](https://arxiv.org) *(coming soon)*

👨‍💻 **Developer:** [Manas Sharma](https://www.linkedin.com/in/manassharma07)

⭐ **Star the repo** if you find it useful!
""")

st.sidebar.markdown("---")

# Installation
with st.sidebar.expander("📦 Installation Instructions"):
    st.code("""
# Install PyFock
pip install pyfock

# For GPU support
pip install cupy-cuda12x  # or appropriate version

# Install dependencies
pip install pyscf numba numpy scipy
""", language="bash")

st.sidebar.markdown("---")

# Performance highlights
with st.sidebar.expander("⚡ Performance Highlights"):
    st.markdown("""
**CPU Performance:**
- Upto 2x faster than PySCF
- Strong scaling up to 32 cores
- ~O(N²·⁰⁵) scaling with basis functions

**GPU Acceleration:**
- Up to **14× speedup** vs 4-core CPU
- Single A100 GPU handles 4000+ basis functions
- Consumer GPUs (RTX series) supported
""")

st.sidebar.markdown("---")
st.sidebar.markdown("*Made with PyFock by PhysWhiz*")
st.sidebar.markdown("*Pure Python • Numba JIT • GPU Ready*")

def strip_ansi(text):
    ansi_escape = re.compile(r'\x1B(?:[@-Z\\-_]|\[[0-?]*[ -/]*[@-~])')
    return ansi_escape.sub('', text)

# Helper: create an ASE Atoms object or fallback
def _parse_xyz_to_atoms(xyz_text):
    # ASE can read from string via io
    return ase_read(StringIO(xyz_text), format='xyz')

# get_structure_viz2 implementation (based on sample provided)
def get_structure_viz2(atoms_obj, style='stick', width=400, height=400):
    xyz_str = ""
    xyz_str += f"{len(atoms_obj)}\n"
    xyz_str += "Structure\n"
    for atom in atoms_obj:
        # atom may be ASE Atoms or fallback MiniAtom
        sym = atom.symbol if hasattr(atom, 'symbol') else atom.get_chemical_symbols()[0]
        pos = atom.position if hasattr(atom, 'position') else atom.position
        xyz_str += f"{sym} {pos[0]:.6f} {pos[1]:.6f} {pos[2]:.6f}\n"
    view = py3Dmol.view(width=width, height=height)
    view.addModel(xyz_str, "xyz")
    if style.lower() == 'ball-stick':
        view.setStyle({'stick': {'radius': 0.2}, 'sphere': {'scale': 0.3}})
    elif style.lower() == 'stick':
        view.setStyle({'stick': {}})
    elif style.lower() == 'ball':
        view.setStyle({'sphere': {'scale': 0.4}})
    else:
        view.setStyle({'stick': {'radius': 0.15}})
    try:
        pbc_any = atoms_obj.pbc.any()
    except Exception:
        pbc_any = False
    
    view.zoomTo()
    view.setBackgroundColor('white')
    return view

# === Cube visualization function ===
def visualize_cube_in_component(cube_content, title, iso_val, opac, width=420, height=360):
    # return the HTML for embedding so it can be used inside columns
    view = py3Dmol.view(width=width, height=height)
    view.addModel(cube_content, 'cube')
    view.setStyle({'sphere': {'colorscheme': 'Jmol', 'scale': 0.3}, 
                    'stick': {'colorscheme': 'Jmol', 'radius': 0.2}})
    
    # For orbitals, show both lobes; for density, show only positive + negative appropriately
    if 'Density' not in title:
        view.addVolumetricData(cube_content, 'cube', 
                                {'isoval': -abs(iso_val), 'color': 'blue', 'opacity': opac})
    view.addVolumetricData(cube_content, 'cube', 
                            {'isoval': abs(iso_val), 'color': 'red', 'opacity': opac})
    
    view.zoomTo()
    view.setClickable({'clickable': 'true'})
    view.enableContextMenu({'contextMenuEnabled': 'true'})
    view.show()
    view.render()
    # don't spin automatically; allow user to toggle via JS later if desired
    t = view.js()
    html_content = f"{t.startjs}{t.endjs}"
    return html_content

# Example XYZ files
EXAMPLE_MOLECULES = {
    "Water": """3
Water molecule
O     0.000000    0.000000    0.117790
H     0.000000    0.755453   -0.471161
H     0.000000   -0.755453   -0.471161""",

    "Acetone": """10
Acetone molecule
O       0.000247197289657     -1.311344859924947      0.000033372829371
C       0.000008761532627     -0.103796835732344      0.000232428229233
C       1.285011287026515      0.689481114475586     -0.000005118102586
C      -1.285310305026895      0.688972899031773     -0.000008308924344
H       1.326033303131908      1.335179621002258     -0.879574382138206
H       1.324106820690737      1.339904097018082      0.876116213728557
H       2.136706543431990      0.014597477886500      0.002551051783708
H      -2.136748467815155      0.013761611436540      0.002550873429523
H      -1.326572603227431      1.334639056170679     -0.879590243114624
H      -1.324682534264540      1.339405821389821      0.876094109485741""",

    "Acetonitrile": """6
Acetonitrile molecule
N       1.238504335855378     -0.000002648443155     -0.000006432209062
C      -1.367114484335133      0.000020833177302      0.000019691866100
C       0.091462639284387      0.000004939432085     -0.000014578846105
H      -1.737603541725548     -0.830506551932429      0.597694569897051
H      -1.737658959168382     -0.102297452169442     -1.018057029309221
H      -1.737589980824850      0.932880879510178      0.420463777423243""",

    "Methane": """5
Methane molecule
C     0.000000    0.000000    0.000000
H     0.629118    0.629118    0.629118
H    -0.629118   -0.629118    0.629118
H    -0.629118    0.629118   -0.629118
H     0.629118   -0.629118   -0.629118""",
    
    "Benzene": """12
Benzene molecule
C     1.395890    0.000000    0.000000
C     0.697945    1.209021    0.000000
C    -0.697945    1.209021    0.000000
C    -1.395890    0.000000    0.000000
C    -0.697945   -1.209021    0.000000
C     0.697945   -1.209021    0.000000
H     2.482610    0.000000    0.000000
H     1.241305    2.149540    0.000000
H    -1.241305    2.149540    0.000000
H    -2.482610    0.000000    0.000000
H    -1.241305   -2.149540    0.000000
H     1.241305   -2.149540    0.000000""",

    "Ammonia": """4
Ammonia molecule
N     0.000000    0.000000    0.100000
H     0.945000    0.000000   -0.266000
H    -0.472500    0.818000   -0.266000
H    -0.472500   -0.818000   -0.266000""",

    "Carbon Dioxide": """3
Carbon dioxide molecule
C     0.000000    0.000000    0.000000
O     0.000000    0.000000    1.160000
O     0.000000    0.000000   -1.160000""",

    "Hydrogen Peroxide": """4
Hydrogen peroxide molecule
O     0.000000    0.000000    0.000000
O     1.450000    0.000000    0.000000
H     0.000000    0.930000    0.000000
H     1.450000   -0.930000    0.000000""",

    "Formaldehyde": """4
Formaldehyde molecule
C     0.000000    0.000000    0.000000
O     1.200000    0.000000    0.000000
H    -0.550000    0.940000    0.000000
H    -0.550000   -0.940000    0.000000""",

    "Hydrogen Cyanide": """3
Hydrogen cyanide molecule
H     0.000000    0.000000    0.000000
C     1.065000    0.000000    0.000000
N     2.232000    0.000000    0.000000""",

    "Acetylene": """4
Acetylene molecule
H     0.000000    0.000000    0.000000
C     0.601000    0.000000    0.000000
C     1.764000    0.000000    0.000000
H     2.365000    0.000000    0.000000""",

    "Ethylene": """6
Ethylene molecule
C     0.000000    0.000000    0.000000
C     1.339000    0.000000    0.000000
H    -0.540000    0.930000    0.000000
H    -0.540000   -0.930000    0.000000
H     1.879000    0.930000    0.000000
H     1.879000   -0.930000    0.000000""",

    "Ethane": """8
Ethane molecule
C     0.000000    0.000000    0.000000
C     1.540000    0.000000    0.000000
H    -0.540000    0.930000    0.000000
H    -0.540000   -0.930000    0.000000
H     0.000000    0.000000    1.090000
H     2.080000    0.930000    0.000000
H     2.080000   -0.930000    0.000000
H     1.540000    0.000000   -1.090000""",

    "Formic Acid": """5
Formic acid molecule
C     0.000000    0.000000    0.000000
O     1.200000    0.000000    0.000000
O    -0.600000    1.100000    0.000000
H     1.700000    0.900000    0.000000
H    -0.600000   -0.900000    0.000000""",

    "Hydrogen Sulfide": """3
Hydrogen sulfide molecule
S     0.000000    0.000000    0.000000
H     0.960000    0.000000    0.000000
H    -0.480000    0.830000    0.000000""",

    "Tetrahydrofuran": """13
Tetrahydrofuran molecule
O       1.216699382773870     -0.000516422279060     -0.000000629115220
C      -1.016993686921708     -0.728737145702576     -0.227053769599122
C      -1.016379073694511      0.729554513122363      0.227039749920194
C       0.395888395971846     -1.160306071902647      0.143905791563541
C       0.396849106350064      1.159943366322044     -0.143953936744947
H      -1.782338327214497     -1.336082444384147      0.254184465508221
H      -1.159922021048819     -0.787714532025205     -1.307880287340100
H      -1.781241112969484      1.337526054547069     -0.254174631443963
H      -1.159223957468540      0.788641683930875      1.307872257115590
H       0.441717113009142     -1.507461993885669      1.181993316376597
H       0.789985254249466     -1.947807118924544     -0.499943162908684
H       0.442962196586526      1.507022288075474     -1.182056728365057
H       0.791596740442480      1.947137848999537      0.499867568609654""",

    "Pyrrole": """10
Pyrrole molecule
N       0.003181105319591     -1.154989666506124      0.000060405177869
C      -1.117737448486888     -0.370847266235924     -0.000011684111207
C       1.119766547077291     -0.364686948222411      0.000018496752116
C      -0.713366217362578      0.937235325591941     -0.000096441964893
C       0.708206650120975      0.941149612298515     -0.000053474487459
H       0.005944376523274     -2.157707636293536     -0.000153625688957
H      -2.102541509509273     -0.805490708525106      0.000110657494436
H       2.106949265704444     -0.793899550618045      0.000097097601888
H      -1.362506178052089      1.796728992178547     -0.000156325233850
H       1.352603398906602      1.804207848609518     -0.000015105544776""",

    "Dimethyl Ether": """9
Dimethyl ether molecule
O       0.000004977280923      0.530020309529034     -0.000005108317779
C       1.164212882530410     -0.260593898872961      0.000001269096639
C      -1.164190191714442     -0.260612616856000     -0.000020896710708
H       1.210503469986643     -0.900685141083738      0.889522583088480
H       2.020207692369726      0.410753932008613      0.000001864695406
H       1.210505099946668     -0.900685096089216     -0.889519959770612
H      -1.210516092616227     -0.900663538561855      0.889526642712922
H      -1.210427870830911     -0.900750324054947     -0.889509722206098
H      -2.020199970977986      0.410716373006091     -0.000096671211569""",

}

# XC functional mapping to libxc codes
XC_FUNCTIONALS = {
    "LDA (SVWN5)": (1, 7),  # Slater exchange + VWN5 correlation
    "PBE": (101, 130),       # PBE exchange + PBE correlation
    "BLYP": (106, 131),      # Becke88 exchange + LYP correlation
    "BP86": (106, 132),      # Becke88 exchange + P86 correlation
}

BASIS_SETS = ["sto-3g", "sto-6g", "3-21G", "4-31G", "6-31G", "6-31+G", "6-31++G", "cc-pvDZ", "def2-SVP", "def2-TZVP"]

# Main title
st.title("⚛️ PyFock GUI - Interactive DFT Calculations")
st.markdown("---")

# Input section
st.header("1. DFT Setup")

col1, col2 = st.columns([1.3, 1])

with col1:
    st.subheader("Molecule Input")
    
    # Molecule selection
    molecule_choice = st.selectbox(
        "Select example molecule or paste custom XYZ:",
        [
            "Water", # correct
            "Acetone", # correct
            "Tetrahydrofuran", # correct
            "Pyrrole", # correct
            "Dimethyl ether", # correct
            # "Acetonitrile", 
            # "Methane",
            "Benzene", # correct
            # "Ammonia",
            "Carbon Dioxide", # correct
            "Hydrogen Peroxide",
            # "Formaldehyde",
            # "Hydrogen Cyanide",
            # "Acetylene",
            # "Ethylene",
            # "Ethane",
            "Formic Acid",
            "Hydrogen Sulfide", # correct
            "Custom"
        ]
    )
    
    if molecule_choice == "Custom":
        xyz_content = st.text_area(
            "Paste XYZ coordinates:",
            height=200,
            placeholder="3\nWater molecule\nO 0.0 0.0 0.0\nH 0.757 0.586 0.0\nH -0.757 0.586 0.0"
        )
    else:
        xyz_content = st.text_area(
            "XYZ coordinates:",
            value=EXAMPLE_MOLECULES[molecule_choice],
            height=200
        )

    # === NEW: Structure visualization right at molecule selection ===
    # This uses ASE if available; otherwise, show a simple py3Dmol view from the XYZ string.
    with col2:
        if xyz_content and xyz_content.strip():
            st.markdown("### Molecule Visualization", unsafe_allow_html=True)
            viz_style = st.selectbox("Select Visualization Style:", ["ball-stick", "stick", "ball"], key="viz_style_select")
            atoms_obj = _parse_xyz_to_atoms(xyz_content)
        
            # Render py3Dmol
            view_3d = get_structure_viz2(atoms_obj, style=viz_style, width=400, height=400)
            # Use components.html to insert the viewer HTML
            try:
                st.components.v1.html(view_3d._make_html(), width=420, height=420)
            except Exception:
                # fallback to js html
                t = view_3d.js()
                html_content = f"{t.startjs}{t.endjs}"
                components.html(html_content, height=420, width=420)

            # Structure information
            st.markdown("### Structure Information")
            atoms_info = {
                "Number of Atoms": len(atoms_obj),
                "Chemical Formula": atoms_obj.get_chemical_formula() if hasattr(atoms_obj, 'get_chemical_formula') else "".join(atoms_obj.get_chemical_symbols()),
                "Atom Types": ", ".join(sorted(list(set(atoms_obj.get_chemical_symbols()))))
            }
            
            for key, value in atoms_info.items():
                st.write(f"**{key}:** {value}")

with col1:
    st.subheader("Calculation Settings")
    
    basis_set = st.selectbox("Basis Set:", BASIS_SETS, index=0)
    auxbasis = st.text_input("Auxiliary Basis:", value="def2-universal-jfit")
    
    xc_functional = st.selectbox(
        "XC Functional:",
        list(XC_FUNCTIONALS.keys()),
        index=0
    )
    
    max_iterations = st.number_input("Max Iterations:", min_value=1, max_value=16, value=14)
    conv_crit = st.number_input("Convergence Criterion:", min_value=1e-7, max_value=1e-3, value=1e-6, format="%.1e")
    ncores = 1#st.number_input("Number of Cores:", min_value=1, max_value=8, value=4)
    use_pyscf_grids = st.checkbox("Use PySCF Grids", value=True, help="Use either PySCF grids for both PyFock and PySCF DFT calculation or PyFock grids for both PyFock and PySCF DFT calculaiton. Using PySCF grids is recommended as those are more efficient and also makes the comparison with PySCF consistent.") 
    compare_pyscf = st.checkbox("Compare with PySCF (may take longer)", value=False, help="Runs a KS-DFT calculation using same settings in PySCF for energy comparison.")

st.markdown("---")

# Visualization settings
st.header("2. Cube Generation and Visualization Settings")
col3, col4, col5 = st.columns(3)

with col3:
    cube_resolution = st.slider("Cube File Resolution (nx=ny=nz):", 30, 70, 40)
with col4:
    isovalue = st.number_input("Isovalue:", 0.0, 1.0, value=0.05, step=0.001, format="%.6f")
with col5:
    opacity = st.slider("Opacity:", 0.0, 1.0, value=0.90, step=0.01)



st.markdown("---")

# Run calculation button
if st.button("🚀 Run DFT Calculation", type="primary"):
    
    # Validate XYZ input
    if not xyz_content.strip():
        st.error("Please provide XYZ coordinates!")
        st.stop()
    
    
    
    # Progress tracking
    progress_bar = st.progress(0)
    status_text = st.empty()
    
    try:
        # Capture stdout/stderr into buffer and show later in the app
        log_buffer = _io.StringIO()
        with contextlib.redirect_stdout(log_buffer), contextlib.redirect_stderr(log_buffer):
            # Setup environment variables
            status_text.text("Setting up environment...")
            progress_bar.progress(5)
            
            os.environ['OMP_NUM_THREADS'] = str(ncores)
            os.environ["OPENBLAS_NUM_THREADS"] = str(ncores)
            os.environ["MKL_NUM_THREADS"] = str(ncores)
            os.environ["VECLIB_MAXIMUM_THREADS"] = str(ncores)
            os.environ["NUMEXPR_NUM_THREADS"] = str(ncores)
            
            # Import PyFock modules
            status_text.text("Importing PyFock modules...")
            progress_bar.progress(10)
            
            from pyfock import Basis, Mol, DFT, Utils, Grids
            
            # Create temporary XYZ file
            status_text.text("Creating molecule object...")
            progress_bar.progress(15)
            
            with tempfile.NamedTemporaryFile(mode='w', suffix='.xyz', delete=False) as f:
                f.write(xyz_content)
                xyz_file = f.name
            
            # Initialize molecule
            mol = Mol(coordfile=xyz_file)
            
            # Initialize basis sets
            status_text.text(f"Loading basis set: {basis_set}...")
            progress_bar.progress(20)
            
            basis = Basis(mol, {'all': Basis.load(mol=mol, basis_name=basis_set)})
            auxbasis_obj = Basis(mol, {'all': Basis.load(mol=mol, basis_name=auxbasis)})
            
            # Check actual basis function count
            n_basis = basis.bfs_nao
            if n_basis > 120:
                st.error(f"❌ This system has {n_basis} basis functions, exceeding the limit of 120. Please use a smaller basis set or fewer atoms.")
                os.unlink(xyz_file)
                st.stop()
            
            st.info(f"✓ System has {n_basis} basis functions (within limit)")
            
            # Get XC functional codes
            funcx, funcc = XC_FUNCTIONALS[xc_functional]
            funcidcrysx = [funcx, funcc]
            funcidpyscf = f"{funcx},{funcc}"
            
            # Generating grids
            status_text.text("Generating numerical grids...")
            progress_bar.progress(22)
            if use_pyscf_grids:
                # PySCF grids
                molPySCF = gto.Mole()
                molPySCF.atom = xyz_file
                molPySCF.basis = basis_set
                molPySCF.cart = True
                molPySCF.verbose = 0
                molPySCF.build()
                grids = gen_grid.Grids(molPySCF)
                grids.level = 3        # optional: quality of grid (0–9 approx)
                grids.prune = None     # disable pruning if you want the full Lebedev mesh
                grids.build(with_non0tab=True)
            else:
                # PyFock grids
                grids = Grids(mol, basis=basis, level = 3, radial_precision=1.0e-13, ncores=ncores)

            

            # Initialize DFT object
            status_text.text("Initializing DFT calculation...")
            progress_bar.progress(25)
            
            dftObj = DFT(mol, basis, auxbasis_obj, xc=funcidcrysx, grids=grids, gridsLevel=3)
            dftObj.conv_crit = conv_crit
            dftObj.max_itr = max_iterations
            dftObj.ncores = 1
            dftObj.save_ao_values = True
            
            # Run SCF calculation
            status_text.text("Running SCF calculation... This may take a few moments.")
            progress_bar.progress(30)
            
            start_time = time.time()
            energyPyFock, dmat = dftObj.scf()
            pyfock_time = time.time() - start_time
            
            progress_bar.progress(50)
            
            # Display results
            st.success(f"✅ PyFock calculation completed in {pyfock_time:.2f} seconds!")
            
            st.header("3. Results")
            
            # Energy and basic properties
            col6, col7, col8 = st.columns(3)
            
            with col6:
                st.metric("Total Energy (PyFock)", f"{energyPyFock:.8f} Ha")
                
                with st.expander("Energy Components"):
                    import pandas as pd
                    energy_df = pd.DataFrame({
                        "Component": [
                            "Kinetic Energy",
                            "Nuclear-Electron Attraction", 
                            "Electron-Electron Repulsion",
                            "Exchange-Correlation",
                            "Nuclear Repulsion"
                        ],
                        "Energy (Ha)": [
                            f"{dftObj.Kinetic_energy:.8f}",
                            f"{dftObj.Nuc_energy:.8f}",
                            f"{dftObj.J_energy:.8f}",
                            f"{dftObj.XC_energy:.8f}",
                            f"{dftObj.Nuclear_repulsion_energy:.8f}"
                        ]
                    })
                    st.dataframe(energy_df, hide_index=True, use_container_width=True)
            
            with col7:
                # Calculate HOMO-LUMO gap
                occupied = np.where(dftObj.mo_occupations > 1e-8)[0]
                if len(occupied) > 0 and len(occupied) < len(dftObj.mo_energies):
                    homo_idx = occupied[-1]
                    lumo_idx = homo_idx + 1
                    homo_energy = dftObj.mo_energies[homo_idx]
                    lumo_energy = dftObj.mo_energies[lumo_idx]
                    gap = (lumo_energy - homo_energy) * 27.2114  # Convert to eV
                    st.metric("HOMO-LUMO Gap", f"{gap:.4f} eV")
                else:
                    homo_idx = None
                    lumo_idx = None
                    st.metric("HOMO-LUMO Gap", "N/A")
            
            with col8:
                # st.metric("SCF Iterations", f"{len(dftObj.energy_list)}")
                st.write('TODO: SCF Iterations')
            
            # MO energies
            st.subheader("Molecular Orbital Energies")
            mo_energies_ev = dftObj.mo_energies * 27.2114  # Convert to eV
            
            col9, col10 = st.columns(2)
            with col9:
                if homo_idx is not None:
                    st.write(f"**HOMO (orbital {homo_idx}):** {mo_energies_ev[homo_idx]:.4f} eV")
            with col10:
                if lumo_idx is not None:
                    st.write(f"**LUMO (orbital {lumo_idx}):** {mo_energies_ev[lumo_idx]:.4f} eV")
            
            # Show MO energies
            with st.expander("View All MO Energies"):
                mo_data = {
                    "Orbital": list(range(len(mo_energies_ev))),
                    "Energy (eV)": [f"{e:.6f}" for e in mo_energies_ev],
                    "Occupation": dftObj.mo_occupations
                }
                st.dataframe(mo_data, height=300)
            # Density matrix expander and download ===
            with st.expander("Density Matrix (dmat) — view / download"):
                try:
                    st.write(dmat)
                    
                except Exception as e:
                    st.write("Failed to show density matrix:", str(e))

            # Generate cube files
            status_text.text("Generating cube files for visualization...")
            progress_bar.progress(60)
            
            cube_files = {}
            
            if homo_idx is not None:
                # HOMO cube
                status_text.text("Generating HOMO cube file...")
                with tempfile.NamedTemporaryFile(mode='w', suffix='_HOMO.cube', delete=False) as f:
                    homo_cube_file = f.name
                
                Utils.write_orbital_cube(
                    mol, basis, dftObj.mo_coefficients[:, homo_idx],
                    homo_cube_file, nx=cube_resolution, ny=cube_resolution, nz=cube_resolution,
                    ncores=ncores
                )
                
                with open(homo_cube_file, 'r') as f:
                    cube_files['HOMO'] = f.read()
                
                progress_bar.progress(70)
            
            if lumo_idx is not None:
                # LUMO cube
                status_text.text("Generating LUMO cube file...")
                with tempfile.NamedTemporaryFile(mode='w', suffix='_LUMO.cube', delete=False) as f:
                    lumo_cube_file = f.name
                
                Utils.write_orbital_cube(
                    mol, basis, dftObj.mo_coefficients[:, lumo_idx],
                    lumo_cube_file, nx=cube_resolution, ny=cube_resolution, nz=cube_resolution,
                    ncores=ncores
                )
                
                with open(lumo_cube_file, 'r') as f:
                    cube_files['LUMO'] = f.read()
                
                progress_bar.progress(80)
            
            # Density cube
            status_text.text("Generating electron density cube file...")
            with tempfile.NamedTemporaryFile(mode='w', suffix='_density.cube', delete=False) as f:
                density_cube_file = f.name
            
            Utils.write_density_cube(
                mol, basis, dftObj.dmat,
                density_cube_file, nx=cube_resolution, ny=cube_resolution, nz=cube_resolution,
                ncores=ncores
            )
            
            with open(density_cube_file, 'r') as f:
                cube_files['Density'] = f.read()
            
            progress_bar.progress(85)
            
            
            
            # Display visualizations side-by-side
            st.subheader("4. Visualizations")
            # create columns for HOMO, LUMO and Density (as available)
            vis_cols = []
            num_vis = len([k for k in cube_files.keys() if cube_files.get(k)])
            if num_vis == 0:
                st.info("No cube visualizations available.")
            else:
                # Arrange into up to three columns side-by-side
                if 'HOMO' in cube_files and 'LUMO' in cube_files and 'Density' in cube_files:
                    c1, c2, c3 = st.columns(3)
                    vis_cols = [c1, c2, c3]
                    mapping = [('HOMO', c1), ('LUMO', c2), ('Density', c3)]
                else:
                    # pack present visualizations into equal columns
                    keys_present = list(cube_files.keys())
                    cols = st.columns(len(keys_present))
                    vis_cols = cols
                    mapping = list(zip(keys_present, cols))
                
                for title, col in mapping:
                    if title in cube_files:
                        with col:
                            st.markdown(f"#### {title}")
                            html_blob = visualize_cube_in_component(cube_files[title], title, isovalue, opacity)
                            components.html(html_blob, height=380, width=420)

            col14, col15, col16 = st.columns(3)
            
            # Helper to create base64 download links (avoids widget-triggered reruns)
            def make_download_link(content, filename, mimetype="text/plain"):
                if isinstance(content, str):
                    b = content.encode()
                else:
                    b = content
                b64 = base64.b64encode(b).decode()
                return f'<a href="data:{mimetype};base64,{b64}" download="{filename}">📥 Download {filename}</a>'
            
            with col14:
                if 'HOMO' in cube_files:
                    st.markdown(make_download_link(cube_files['HOMO'], "homo.cube"), unsafe_allow_html=True)
            
            with col15:
                if 'LUMO' in cube_files:
                    st.markdown(make_download_link(cube_files['LUMO'], "lumo.cube"), unsafe_allow_html=True)
            
            with col16:
                if 'Density' in cube_files:
                    st.markdown(make_download_link(cube_files['Density'], "density.cube"), unsafe_allow_html=True)
            
            progress_bar.progress(90)
            
            # === Allow visualizing any orbital ===
            @st.fragment
            def func_viz_any_mo():
                st.subheader("Visualize any MO")
                mo_idx_choice = None
                if hasattr(dftObj, 'mo_energies'):
                    max_orb = len(dftObj.mo_energies) - 1
                    # Show a slider/selectbox to pick orbital
                    mo_idx_choice = st.number_input("Select orbital index to visualize:", min_value=0, max_value=max_orb, value=homo_idx if homo_idx is not None else 0, step=1)
                    # if st.button("Generate and Show Selected MO", key="gen_orb_btn"):
                    status_text.text(f"Generating cube for MO index {mo_idx_choice} ...")
                    with tempfile.NamedTemporaryFile(mode='w', suffix=f'_MO{mo_idx_choice}.cube', delete=False) as f:
                        mo_cube_file = f.name
                    Utils.write_orbital_cube(
                        mol, basis, dftObj.mo_coefficients[:, int(mo_idx_choice)],
                        mo_cube_file, nx=cube_resolution, ny=cube_resolution, nz=cube_resolution,
                        ncores=ncores
                    )
                    with open(mo_cube_file, 'r') as f:
                        cube_files[f"MO_{mo_idx_choice}"] = f.read()
                    # show it in a small area
                    html_blob = visualize_cube_in_component(cube_files[f"MO_{mo_idx_choice}"], f"MO {mo_idx_choice}", isovalue, opacity)
                    st.markdown(f"#### MO {mo_idx_choice}")
                    components.html(html_blob, height=380, width=420)
            
            func_viz_any_mo()
                
            
            # PySCF comparison
            if compare_pyscf:
                status_text.text("Running PySCF calculation for comparison...")
                
                try:
                    if not use_pyscf_grids:
                        molPySCF = gto.Mole()
                        molPySCF.atom = xyz_file
                        molPySCF.basis = basis_set
                        molPySCF.cart = True
                        molPySCF.verbose = 5
                        molPySCF.build()
                    
                    mf = dft.rks.RKS(molPySCF).density_fit(auxbasis=auxbasis)
                    mf.xc = funcidpyscf
                    # mf.direct_scf = False
                    mf.conv_tol = conv_crit
                    mf.max_cycle = max_iterations
                    # mf.grids.level = 5
                    mf.grids.coords = grids.coords
                    mf.grids.weights = grids.weights

                    # Disable PySCF's automatic grid generation
                    mf.grids.build = lambda *args, **kwargs: None
                    
                    start_pyscf = time.time()
                    energyPySCF = mf.kernel(dm0 = mf.init_guess_by_1e(molPySCF))
                    pyscf_time = time.time() - start_pyscf
                    
                    st.subheader("5. Comparison with PySCF")
                    
                    col11, col12, col13 = st.columns(3)
                    
                    with col11:
                        st.metric("PySCF Energy", f"{energyPySCF:.8f} Ha")
                    
                    with col12:
                        energy_diff = abs(energyPyFock - energyPySCF) * 1000  # in mHa
                        st.metric("Energy Difference", f"{energy_diff:.6f} mHa")
                    
                    with col13:
                        speedup = pyscf_time / pyfock_time
                        # st.metric("PyFock Speedup", f"{speedup:.2f}x" if speedup > 1 else f"{1/speedup:.2f}x slower")
                
                        # st.write(f"**PyFock time:** {pyfock_time:.2f} s | **PySCF time:** {pyscf_time:.2f} s")
                    
                except Exception as e:
                    st.warning(f"PySCF comparison failed: {str(e)}")
            
            progress_bar.progress(95)
            
            # Downloads section
            st.subheader("6. INPUT Script Generation")
            
            
            
            
            
            
            
            # Generate Python script
            status_text.text("Generating Python script...")
            
            python_script = """# PyFock DFT Calculation Script
# Generated by PyFock GUI

import os
ncores = {ncores}
os.environ['OMP_NUM_THREADS'] = str(ncores)
os.environ["OPENBLAS_NUM_THREADS"] = str(ncores)
os.environ["MKL_NUM_THREADS"] = str(ncores)
os.environ["VECLIB_MAXIMUM_THREADS"] = str(ncores)
os.environ["NUMEXPR_NUM_THREADS"] = str(ncores)

from pyfock import Basis, Mol, DFT, Utils
import numpy as np

# XYZ coordinates
xyz_content = \"\"\"
{xyz_content}
\"\"\"

# Save XYZ to file
with open('molecule.xyz', 'w') as f:
    f.write(xyz_content)

# Calculation parameters
basis_set_name = '{basis_set}'
auxbasis_name = '{auxbasis}'
funcx = {funcx}
funcc = {funcc}
funcidcrysx = [funcx, funcc]

# Initialize molecule and basis
mol = Mol(coordfile='molecule.xyz')
basis = Basis(mol, {{'all': Basis.load(mol=mol, basis_name=basis_set_name)}})
auxbasis = Basis(mol, {{'all': Basis.load(mol=mol, basis_name=auxbasis_name)}})

# Setup DFT calculation
dftObj = DFT(mol, basis, auxbasis, xc=funcidcrysx)
dftObj.conv_crit = {conv_crit}
dftObj.max_itr = {max_iterations}
dftObj.ncores = ncores
dftObj.save_ao_values = True

# Run SCF
energy, dmat = dftObj.scf()

print(f"Total Energy: {{energy:.8f}} Ha")

# Find HOMO and LUMO
occupied = np.where(dftObj.mo_occupations > 1e-8)[0]
if len(occupied) > 0 and len(occupied) < len(dftObj.mo_energies):
    homo_idx = occupied[-1]
    lumo_idx = homo_idx + 1
    
    # Generate cube files
    Utils.write_orbital_cube(mol, basis, dftObj.mo_coefficients[:, homo_idx], 
                            'HOMO.cube', nx={cube_resolution}, ny={cube_resolution}, 
                            nz={cube_resolution}, ncores=ncores)
    
    Utils.write_orbital_cube(mol, basis, dftObj.mo_coefficients[:, lumo_idx], 
                            'LUMO.cube', nx={cube_resolution}, ny={cube_resolution}, 
                            nz={cube_resolution}, ncores=ncores)

# Generate density cube
Utils.write_density_cube(mol, basis, dftObj.dmat, 'density.cube', 
                        nx={cube_resolution}, ny={cube_resolution}, 
                        nz={cube_resolution}, ncores=ncores)

print("Cube files generated successfully!")
"""
            parameters = {
                'ncores': ncores,
                'xyz_content': xyz_content,
                'basis_set': basis_set,
                'auxbasis': auxbasis, # Often a different basis set for auxiliary functions
                'funcx': funcx,        # Example ID for exchange functional (like LDA_X)
                'funcc': funcc,        # Example ID for correlation functional (like LDA_C_PZ)
                'conv_crit': conv_crit,
                'max_iterations': max_iterations,
                'cube_resolution': cube_resolution
            }
            python_script = python_script.format(**parameters)
            with st.expander("Input Script to Run the Above PyFock Calculation"):
                st.code(python_script)
            # Provide script as base64 link (no rerun on click)
            st.markdown(make_download_link(python_script, "pyfock_calculation.py", mimetype="text/x-python"), unsafe_allow_html=True)
            
            progress_bar.progress(100)
            status_text.text("✅ All tasks completed!")
            
            # === Show captured stdout/stderr logs ===
            st.subheader("7. Calculation Log Output")
            log_buffer.seek(0)
            log_text = log_buffer.read()
            if log_text.strip():
                st.code(strip_ansi(log_text))
            else:
                st.write("No log output was captured.")
            
            
            # Cleanup
            os.unlink(xyz_file)
            if 'homo_cube_file' in locals():
                os.unlink(homo_cube_file)
            if 'lumo_cube_file' in locals():
                os.unlink(lumo_cube_file)
            if 'density_cube_file' in locals():
                os.unlink(density_cube_file)
            if 'mo_cube_file' in locals():
                try:
                    os.unlink(mo_cube_file)
                except Exception:
                    pass

    except ImportError as e:
        st.error(f"❌ Import Error: {str(e)}")
        st.info("Make sure PyFock is installed: `pip install pyfock`")
        progress_bar.empty()
        status_text.empty()
    except Exception as e:
        st.error(f"❌ Calculation failed: {str(e)}")
        st.info("Please check your input parameters and try again.")
        progress_bar.empty()
        status_text.empty()
        
        # Cleanup on error
        if 'xyz_file' in locals():
            try:
                os.unlink(xyz_file)
            except:
                pass

else:
    # Initial instructions
    st.info("👆 Configure your calculation parameters above and click 'Run DFT Calculation' to start!")
    
    st.markdown("""
    ### Getting Started
    
    1. **Choose a molecule**: Select from examples or paste your own XYZ coordinates
    2. **Select calculation parameters**: Basis set, functional, convergence criteria
    3. **Adjust visualization settings**: Cube resolution, isovalue, opacity
    4. **Run calculation**: Click the button above
    5. **Explore results**: View energies, MO properties, and 3D visualizations
    6. **Download**: Get cube files and a Python script to reproduce the calculation
    
    ### Notes
    
    - Calculations are limited to ~120 basis functions for Streamlit Cloud
    - Smaller molecules and basis sets will run faster
    - PySCF comparison adds computation time but validates results
    - All calculations use density fitting for efficiency
    
    ### Example Systems
    
    - **Water**: Quick test system (7 basis functions with sto-3g)
    - **Methane**: Small hydrocarbon (9 basis functions with sto-3g)
    - **Benzene**: Aromatic system (42 basis functions with sto-3g)
    """)

# Footer
st.markdown("---")
st.markdown("""
<div style='text-align: center'>
    <p>PyFock GUI - Pure Python DFT with Numba JIT acceleration</p>
    <p>⚡ Fast • 🎯 Accurate • 🐍 Pure Python</p>
</div>
""", unsafe_allow_html=True)