import os
import json
import gradio as gr
import groq
import numpy as np
import matplotlib.pyplot as plt
import py3Dmol
from pymatgen.core.structure import Structure
from pymatgen.io.cif import CifWriter
from pymatgen.io.xyz import XYZ
from pymatgen.symmetry.analyzer import SpacegroupAnalyzer
import requests
from dotenv import load_dotenv
import tempfile
import base64

# Load environment variables
load_dotenv()

# Retrieve the API key from the environment variable
groq_api_key = os.getenv("GROQ_API_KEY")
if not groq_api_key:
    raise ValueError("GROQ_API_KEY environment variable not set")

# Initialize Groq client - using a simpler initialization to avoid proxies parameter issue
client = groq.Client(api_key=groq_api_key)

# Function to generate crystal structures
def generate_crystal_structure(material_name):
    """Generate a crystal structure for a given material."""
    # This is a simplified version - in a real application, you would use a database or API
    # to get real crystal structure data for materials
    
    # Dictionary of common materials and their crystal structures
    material_structures = {
        "Silicon": {
            "spacegroup": "Fd-3m",
            "lattice": [[0, 5.431/2, 5.431/2], [5.431/2, 0, 5.431/2], [5.431/2, 5.431/2, 0]],
            "species": ["Si"],
            "coords": [[0, 0, 0]]
        },
        "Titanium Dioxide": {
            "spacegroup": "P42/mnm",
            "lattice": [[4.594, 0, 0], [0, 4.594, 0], [0, 0, 2.959]],
            "species": ["Ti", "O", "O"],
            "coords": [[0, 0, 0], [0.3053, 0.3053, 0], [0.3053, 0.6947, 0.5]]
        },
        "Graphene": {
            "spacegroup": "P6/mmm",
            "lattice": [[2.46, 0, 0], [-2.46/2, 2.46*np.sqrt(3)/2, 0], [0, 0, 15]],
            "species": ["C", "C"],
            "coords": [[0, 0, 0], [1/3, 2/3, 0]]
        },
        "Copper": {
            "spacegroup": "Fm-3m",
            "lattice": [[3.615, 0, 0], [0, 3.615, 0], [0, 0, 3.615]],
            "species": ["Cu"],
            "coords": [[0, 0, 0]]
        },
        "Aluminum": {
            "spacegroup": "Fm-3m",
            "lattice": [[4.05, 0, 0], [0, 4.05, 0], [0, 0, 4.05]],
            "species": ["Al"],
            "coords": [[0, 0, 0]]
        },
        "Gold": {
            "spacegroup": "Fm-3m",
            "lattice": [[4.078, 0, 0], [0, 4.078, 0], [0, 0, 4.078]],
            "species": ["Au"],
            "coords": [[0, 0, 0]]
        },
        "Diamond": {
            "spacegroup": "Fd-3m",
            "lattice": [[0, 3.567/2, 3.567/2], [3.567/2, 0, 3.567/2], [3.567/2, 3.567/2, 0]],
            "species": ["C"],
            "coords": [[0, 0, 0]]
        },
        "Graphite": {
            "spacegroup": "P63/mmc",
            "lattice": [[2.46, 0, 0], [-1.23, 2.13, 0], [0, 0, 6.71]],
            "species": ["C", "C", "C", "C"],
            "coords": [[0, 0, 0], [0, 0, 0.5], [1/3, 2/3, 0], [2/3, 1/3, 0.5]]
        }
    }
    
    # Try to match the material name with our database (case insensitive)
    for known_material, structure_data in material_structures.items():
        if material_name.lower() in known_material.lower():
            structure = Structure.from_spacegroup(
                structure_data["spacegroup"],
                lattice=structure_data["lattice"],
                species=structure_data["species"],
                coords=structure_data["coords"]
            )
            return structure
    
    # If material not found, create a generic structure
    return Structure.from_spacegroup(
        "Pm-3m",
        lattice=[[4.0, 0, 0], [0, 4.0, 0], [0, 0, 4.0]],
        species=["X"],
        coords=[[0, 0, 0]]
    )

# Function to get material recommendations from LLM
def get_material_recommendations(query):
    """Get material recommendations from the LLM based on user query."""
    system_prompt = """You are a materials science expert. Your task is to recommend the 3 best materials for a specific application or with certain properties based on the user's query.
    
    For each material, provide:
    1. Material name
    2. Chemical formula
    3. Key properties relevant to the application
    4. Why it's suitable for the application
    5. Any limitations or considerations
    
    Format your response as a JSON object with the following structure:
    {
        "materials": [
            {
                "name": "Material Name",
                "formula": "Chemical Formula",
                "properties": "Key properties relevant to the application",
                "suitability": "Why it's suitable for the application",
                "limitations": "Any limitations or considerations"
            },
            // Second material
            // Third material
        ]
    }
    
    Ensure your response is strictly in this JSON format with no additional text."""
    
    try:
        # Updated to use the chat.completions API with the correct client method
        completion = client.chat.completions.create(
            model="deepseek-r1",
            messages=[
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": query}
            ],
            temperature=0.2,
            max_tokens=1000
        )
        
        response_text = completion.choices[0].message.content
        
        # Extract JSON from the response
        try:
            # Try to parse the entire response as JSON
            recommendations = json.loads(response_text)
        except json.JSONDecodeError:
            # If that fails, try to extract JSON using string manipulation
            json_start = response_text.find('{')
            json_end = response_text.rfind('}') + 1
            if json_start >= 0 and json_end > json_start:
                json_str = response_text[json_start:json_end]
                recommendations = json.loads(json_str)
            else:
                raise ValueError("Could not extract valid JSON from LLM response")
        
        return recommendations
    except Exception as e:
        return {"error": str(e)}

# Function to get crystal structure information
def get_crystal_structure_info(material_name):
    """Get crystal structure information for a given material."""
    try:
        # Generate the crystal structure
        structure = generate_crystal_structure(material_name)
        
        # Create a CIF file
        cif_writer = CifWriter(structure)
        with tempfile.NamedTemporaryFile(suffix='.cif', delete=False) as temp_cif:
            cif_writer.write_file(temp_cif.name)
            cif_path = temp_cif.name
        
        with open(cif_path, 'r') as f:
            cif_content = f.read()
        
        # Create an XYZ file
        with tempfile.NamedTemporaryFile(suffix='.xyz', delete=False) as temp_xyz:
            xyz_path = temp_xyz.name
        
        # Convert structure to XYZ format
        ase_atoms = structure.to_ase()
        from ase.io import write as ase_write
        ase_write(xyz_path, ase_atoms, format='xyz')
        
        with open(xyz_path, 'r') as f:
            xyz_content = f.read()
        
        # Get space group information
        analyzer = SpacegroupAnalyzer(structure)
        spacegroup = analyzer.get_space_group_symbol()
        
        # Generate 3D visualization
        view = py3Dmol.view(width=500, height=400)
        view.addModel(xyz_content, 'xyz')
        view.setStyle({'sphere': {'colorscheme': 'Jmol', 'scale': 0.3},
                       'stick': {'radius': 0.2}})
        view.zoomTo()
        view.spin(True)
        view.setBackgroundColor('white')
        view.render()
        
        # Convert the view to HTML
        html_str = view._make_html()
        
        # Clean up temporary files
        os.unlink(cif_path)
        os.unlink(xyz_path)
        
        return {
            "material_name": material_name,
            "formula": structure.composition.reduced_formula,
            "space_group": spacegroup,
            "num_atoms": len(structure),
            "lattice_parameters": {
                "a": structure.lattice.a,
                "b": structure.lattice.b,
                "c": structure.lattice.c,
                "alpha": structure.lattice.alpha,
                "beta": structure.lattice.beta,
                "gamma": structure.lattice.gamma
            },
            "cif_content": cif_content,
            "xyz_content": xyz_content,
            "visualization_html": html_str
        }
    except Exception as e:
        return {"error": str(e)}

# Function to create a downloadable file
def create_downloadable_file(content, filename):
    """Create a downloadable file with the given content."""
    with open(filename, 'w') as f:
        f.write(content)
    return filename

# Gradio interface
def process_query(query):
    """Process the user query and return material recommendations and crystal structure visualization."""
    try:
        # Get material recommendations from LLM
        recommendations = get_material_recommendations(query)
        
        if "error" in recommendations:
            return f"Error getting recommendations: {recommendations['error']}", None, None, None
        
        # Format the recommendations as text
        recommendation_text = "# Material Recommendations\n\n"
        
        for i, material in enumerate(recommendations.get("materials", [])):
            recommendation_text += f"## {i+1}. {material.get('name', 'Unknown')}\n"
            recommendation_text += f"**Formula:** {material.get('formula', 'N/A')}\n\n"
            recommendation_text += f"**Properties:** {material.get('properties', 'N/A')}\n\n"
            recommendation_text += f"**Suitability:** {material.get('suitability', 'N/A')}\n\n"
            recommendation_text += f"**Limitations:** {material.get('limitations', 'N/A')}\n\n"
        
        # Get crystal structure information for the first recommended material
        cif_file = None
        xyz_file = None
        
        if recommendations.get("materials") and len(recommendations.get("materials")) > 0:
            first_material = recommendations["materials"][0]["name"]
            structure_info = get_crystal_structure_info(first_material)
            
            if "error" in structure_info:
                structure_html = f"<p>Error getting crystal structure: {structure_info['error']}</p>"
            else:
                # Add crystal structure information to the recommendation text
                recommendation_text += f"# Crystal Structure of {structure_info['material_name']}\n\n"
                recommendation_text += f"**Formula:** {structure_info['formula']}\n\n"
                recommendation_text += f"**Space Group:** {structure_info['space_group']}\n\n"
                recommendation_text += f"**Number of Atoms:** {structure_info['num_atoms']}\n\n"
                recommendation_text += "**Lattice Parameters:**\n"
                recommendation_text += f"a = {structure_info['lattice_parameters']['a']:.4f} Å, "
                recommendation_text += f"b = {structure_info['lattice_parameters']['b']:.4f} Å, "
                recommendation_text += f"c = {structure_info['lattice_parameters']['c']:.4f} Å\n"
                recommendation_text += f"α = {structure_info['lattice_parameters']['alpha']:.2f}°, "
                recommendation_text += f"β = {structure_info['lattice_parameters']['beta']:.2f}°, "
                recommendation_text += f"γ = {structure_info['lattice_parameters']['gamma']:.2f}°\n\n"
                
                # Create HTML for the 3D visualization
                structure_html = structure_info['visualization_html']
                
                # Create downloadable files
                cif_file = create_downloadable_file(structure_info['cif_content'], f"{structure_info['material_name'].replace(' ', '_')}.cif")
                xyz_file = create_downloadable_file(structure_info['xyz_content'], f"{structure_info['material_name'].replace(' ', '_')}.xyz")
        else:
            structure_html = "<p>No materials recommended to visualize.</p>"
        
        return recommendation_text, structure_html, cif_file, xyz_file
    except Exception as e:
        return f"Error processing query: {str(e)}", None, None, None

# Create the Gradio interface
with gr.Blocks(title="Material Science Expert") as demo:
    gr.Markdown("# Material Science Expert")
    gr.Markdown("Ask for a material for a specific application or with certain properties.")
    
    with gr.Row():
        with gr.Column():
            query_input = gr.Textbox(
                label="Your Query",
                placeholder="I need a material with high thermal conductivity for electronics cooling.",
                lines=3
            )
            submit_btn = gr.Button("Get Recommendations")
        
    with gr.Row():
        with gr.Column(scale=1):
            recommendations_output = gr.Markdown(label="Material Recommendations")
        with gr.Column(scale=1):
            structure_output = gr.HTML(label="Crystal Structure Visualization")
    
    with gr.Row():
        with gr.Column():
            cif_file_output = gr.File(label="Download CIF File")
            xyz_file_output = gr.File(label="Download XYZ File")
    
    submit_btn.click(
        fn=process_query,
        inputs=[query_input],
        outputs=[recommendations_output, structure_output, cif_file_output, xyz_file_output]
    )
    
    gr.Markdown("""
    ## Example Queries:
    - I need a material with high thermal conductivity for electronics cooling.
    - What are the best materials for solar cell applications?
    - Recommend materials with high strength-to-weight ratio for aerospace applications.
    - I need a transparent conductive material for touchscreens.
    """)

# Launch the app
if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0")