"""
Quantum ESPRESSO MCP Server - Gradio Demo

A demonstration interface for the QE-MCP server that enables
LLMs to run DFT calculations with natural language.
"""

import gradio as gr
import json

# Demo data (since we can't run Docker on HuggingFace Spaces)
DEMO_RESULTS = {
    "Si": {
        "scf": {
            "success": True,
            "total_energy_eV": -214.4906,
            "fermi_energy_eV": 6.2975,
            "converged": True,
            "n_iterations": 4,
            "parameters_used": {"spin_polarized": False, "smearing": "cold", "degauss": 0.02}
        },
        "bandstructure": {
            "success": True,
            "band_gap_eV": 0.56,
            "is_direct": False,
            "vbm_location": "Γ",
            "cbm_location": "X",
            "fermi_energy_eV": 6.2975
        }
    },
    "Fe": {
        "scf": {
            "success": True,
            "total_energy_eV": -3220.5287,
            "fermi_energy_eV": 17.8432,
            "total_magnetization": 7.63,
            "converged": True,
            "n_iterations": 12,
            "parameters_used": {"spin_polarized": True, "smearing": "cold", "degauss": 0.02}
        }
    },
    "Cu": {
        "scf": {
            "success": True,
            "total_energy_eV": -1653.2341,
            "fermi_energy_eV": 12.4521,
            "converged": True,
            "n_iterations": 6,
            "parameters_used": {"spin_polarized": False, "smearing": "cold", "degauss": 0.02}
        }
    },
    "GaAs": {
        "scf": {
            "success": True,
            "total_energy_eV": -312.8765,
            "fermi_energy_eV": 5.1234,
            "converged": True,
            "n_iterations": 5,
            "parameters_used": {"spin_polarized": False, "smearing": "cold", "degauss": 0.02}
        },
        "bandstructure": {
            "success": True,
            "band_gap_eV": 0.48,
            "is_direct": True,
            "vbm_location": "Γ",
            "cbm_location": "Γ",
            "fermi_energy_eV": 5.1234
        }
    }
}

AVAILABLE_ELEMENTS = [
    "Ag", "Al", "Ar", "As", "Au", "B", "Ba", "Be", "Bi", "Br", "C", "Ca", "Cd", "Cl", 
    "Co", "Cr", "Cs", "Cu", "F", "Fe", "Ga", "Ge", "H", "He", "Hf", "Hg", "I", "In", 
    "Ir", "K", "Kr", "La", "Li", "Mg", "Mn", "Mo", "N", "Na", "Nb", "Ne", "Ni", "O", 
    "Os", "P", "Pb", "Pd", "Pt", "Rb", "Re", "Rh", "Ru", "S", "Sb", "Sc", "Se", "Si", 
    "Sn", "Sr", "Ta", "Tc", "Te", "Ti", "Tl", "V", "W", "Xe", "Y", "Zn", "Zr"
]

MCP_TOOLS = """
## 🔧 Available MCP Tools

| Tool | Description |
|------|-------------|
| `qe_run_scf` | Self-consistent field calculation (total energy, Fermi level) |
| `qe_run_relax` | Optimize atomic positions |
| `qe_run_vc_relax` | Variable-cell relaxation (optimize positions AND cell) |
| `qe_workflow_bandstructure` | Complete band structure workflow |
| `qe_workflow_dos` | Density of states calculation |
| `qe_workflow_relax_and_scf` | Relax structure then accurate SCF |
| `qe_load_structure` | Load and inspect atomic structures |
| `qe_get_kpath` | Get high-symmetry k-path for band structure |
| `qe_suggest_kpoints` | Suggest k-point grid based on cell size |
| `qe_list_pseudopotentials` | List available elements (69 total) |
| `qe_validate_structure` | Validate structure and check for issues |
| `qe_status` | Get QE MCP server status |
"""

def run_scf_demo(material: str) -> str:
    """Simulate SCF calculation"""
    material = material.strip()
    if material in DEMO_RESULTS:
        result = DEMO_RESULTS[material]["scf"]
        output = f"""## ⚡ SCF Calculation: {material}

✅ **Success**: {result['success']}
🔋 **Total Energy**: {result['total_energy_eV']:.4f} eV
📊 **Fermi Energy**: {result['fermi_energy_eV']:.4f} eV
🔄 **Converged**: {result['converged']} ({result['n_iterations']} iterations)
"""
        if result.get('total_magnetization'):
            output += f"🧲 **Magnetization**: {result['total_magnetization']:.2f} μB\n"
        output += f"\n⚙️ **Auto-detected parameters**: {json.dumps(result['parameters_used'])}"
        return output
    else:
        return f"""## ⚡ SCF Calculation: {material}

This is a **demo** showing the MCP tool interface.

In the full version, calling `qe_run_scf(structure='{material}')` would:
1. Build the crystal structure using ASE
2. Generate QE input files
3. Run pw.x in Docker container
4. Parse and return results

**Supported elements**: {', '.join(AVAILABLE_ELEMENTS)}
"""

def run_bandstructure_demo(material: str) -> str:
    """Simulate band structure calculation"""
    material = material.strip()
    if material in DEMO_RESULTS and "bandstructure" in DEMO_RESULTS[material]:
        result = DEMO_RESULTS[material]["bandstructure"]
        gap_type = "direct" if result['is_direct'] else "indirect"
        return f"""## 📈 Band Structure: {material}

✅ **Success**: {result['success']}
🎯 **Band Gap**: {result['band_gap_eV']:.2f} eV ({gap_type})
📍 **VBM Location**: {result['vbm_location']}
📍 **CBM Location**: {result['cbm_location']}
📊 **Fermi Energy**: {result['fermi_energy_eV']:.4f} eV

### Interpretation
{"This is a **semiconductor** with a " + gap_type + " band gap." if result['band_gap_eV'] > 0 else "This is a **metal**."}
"""
    else:
        return f"""## 📈 Band Structure: {material}

This is a **demo** showing the MCP tool interface.

In the full version, calling `qe_workflow_bandstructure(structure='{material}')` would:
1. Run SCF calculation
2. Get high-symmetry k-path (Γ-X-W-K-Γ-L-U-W-L-K)
3. Calculate bands along path
4. Analyze band gap

**Try**: Si, GaAs (have demo results)
"""

def show_mcp_config() -> str:
    """Show MCP configuration for Claude Desktop"""
    return """## 🔧 Claude Desktop Configuration

Add this to `~/Library/Application Support/Claude/claude_desktop_config.json` (macOS):

```json
{
  "mcpServers": {
    "quantum-espresso": {
      "command": "uv",
      "args": ["--directory", "/path/to/QE_MCP", "run", "qe-mcp"]
    }
  }
}
```

Then restart Claude Desktop and you can say:
- *"Calculate the total energy of silicon"*
- *"What's the band gap of GaAs?"*
- *"Run a spin-polarized calculation for iron"*

## 📋 Requirements

- Docker (with `qe-local` image)
- Python 3.10+
- uv package manager
"""

def list_elements() -> str:
    """List all available elements"""
    elements_grid = ""
    for i, elem in enumerate(AVAILABLE_ELEMENTS):
        elements_grid += f"`{elem}` "
        if (i + 1) % 10 == 0:
            elements_grid += "\n"
    
    return f"""## 🧪 Supported Elements (69 total)

SG15 ONCV Pseudopotential Library:

{elements_grid}

### Magnetic Elements (auto spin-polarized)
`Fe` `Co` `Ni` `Mn` `Cr` `V` `Gd` `Eu` `Tb` `Dy` `Ho` `Er`
"""

# Create Gradio Interface
with gr.Blocks(
    title="⚛️ Quantum ESPRESSO MCP Server",
    theme=gr.themes.Soft(primary_hue="blue", secondary_hue="purple"),
    css="""
    .gradio-container { max-width: 1200px !important; }
    .tool-card { border: 1px solid #e0e0e0; border-radius: 8px; padding: 16px; margin: 8px 0; }
    """
) as demo:
    gr.Markdown("""
    # ⚛️ Quantum ESPRESSO MCP Server
    
    > **Run DFT calculations with natural language!** An MCP server that enables LLMs to perform 
    > first-principles quantum mechanical simulations.
    
    [![MCP Compatible](https://img.shields.io/badge/MCP-Compatible-blue)](https://modelcontextprotocol.io)
    [![Quantum ESPRESSO](https://img.shields.io/badge/QE-v6.7-green)](https://www.quantum-espresso.org/)
    
    ⚠️ **Note**: This is a demo interface. The full MCP server runs locally with Docker.
    """)
    
    with gr.Tabs():
        with gr.Tab("🧪 Try It"):
            gr.Markdown("### Simulate MCP Tool Calls")
            
            with gr.Row():
                with gr.Column():
                    material_input = gr.Textbox(
                        label="Material Formula",
                        placeholder="Si, Fe, Cu, GaAs...",
                        value="Si"
                    )
                    with gr.Row():
                        scf_btn = gr.Button("⚡ Run SCF", variant="primary")
                        band_btn = gr.Button("📈 Band Structure", variant="secondary")
                
                with gr.Column():
                    output = gr.Markdown(label="Result")
            
            scf_btn.click(run_scf_demo, inputs=[material_input], outputs=[output])
            band_btn.click(run_bandstructure_demo, inputs=[material_input], outputs=[output])
            
            gr.Markdown("**Demo materials**: Si, Fe, Cu, GaAs")
        
        with gr.Tab("🔧 MCP Tools"):
            gr.Markdown(MCP_TOOLS)
            
            gr.Markdown("""
            ### 🎯 8 Prompts for Guided Workflows
            
            | Prompt | Description |
            |--------|-------------|
            | `band_structure` | Calculate electronic band structure |
            | `dos_calculation` | Density of states workflow |
            | `geometry_optimization` | Structure relaxation steps |
            | `convergence_test` | Parameter convergence testing |
            | `surface_calculation` | Surface energy calculations |
            | `magnetic_calculation` | Magnetic properties (Fe, Ni, Co) |
            | `troubleshoot` | Diagnose calculation problems |
            | `compare_structures` | Compare multiple structures |
            """)
        
        with gr.Tab("⚙️ Setup"):
            config_output = gr.Markdown(value=show_mcp_config())
        
        with gr.Tab("🧪 Elements"):
            elements_output = gr.Markdown(value=list_elements())
    
    gr.Markdown("""
    ---
    
    ### 🏗️ Architecture
    
    ```
    User (natural language) → LLM (Claude/GPT) → MCP Protocol → QE-MCP Server → Docker (QE v6.7) → Results
    ```
    
    ### 🛠️ Tech Stack
    - **Quantum ESPRESSO v6.7MaX** - DFT engine
    - **MCP SDK** (mcp>=1.0.0) - Model Context Protocol  
    - **ASE 3.26** - Atomic Simulation Environment
    - **Docker** - Containerized QE
    - **SG15 ONCV** - 69 element pseudopotentials
    
    ---
    
    *Built for [MCP's 1st Birthday Hackathon](https://huggingface.co/MCP-1st-Birthday) 🎂 by [@frimpsjoe](https://huggingface.co/frimpsjoe)*
    """)

if __name__ == "__main__":
    demo.launch()