import os
from pymatgen.ext.matproj import MPRester
import crystal_toolkit.components as ctc
from crystal_toolkit.settings import SETTINGS

import dash
from dash import html, dcc
from dash.dependencies import Input, Output, State

# Set your Materials Project API key
MATERIALS_PROJECT_API_KEY = os.getenv('MATERIALS_PROJECT_API_KEY')

# Initialize the Dash app
app = dash.Dash(__name__, assets_folder=SETTINGS.ASSETS_PATH)
server = app.server  # Expose the server for deployment

# Define the app layout
layout = html.Div([
    dcc.Markdown("## Interactive Crystal Viewer"),
    html.Div([
        html.Div([
            html.Label("Search by Chemical System (e.g., 'Ac-Cd-Ge')"),
            dcc.Input(
                id='query-input',
                type='text',
                value='Ac-Cd-Ge',
                placeholder='Ac-Cd-Ge',
                style={'width': '100%'}
            ),
        ], style={'width': '70%', 'display': 'inline-block', 'verticalAlign': 'top'}),
        html.Div([
            html.Button('Search', id='search-button', n_clicks=0),
        ], style={'width': '28%', 'display': 'inline-block', 'paddingLeft': '2%', 'verticalAlign': 'top'}),
    ], style={'margin-bottom': '20px'}),
    html.Div([
        html.Label("Select Material"),
        dcc.Dropdown(
            id='material-dropdown',
            options=[],  # Empty options initially
            value=None
        ),
    ], style={'margin-bottom': '20px'}),
    html.Button('Display Material', id='display-button', n_clicks=0),
    html.Div([
        html.Div(id='structure-container', style={'width': '48%', 'display': 'inline-block', 'verticalAlign': 'top'}),
        html.Div(id='properties-container',
                 style={'width': '48%', 'display': 'inline-block', 'paddingLeft': '4%', 'verticalAlign': 'top'}),
    ], style={'margin-top': '20px'}),
])


# Function to search for materials
def search_materials(query):
    with MPRester(MATERIALS_PROJECT_API_KEY) as mpr:
        results = mpr.summary.search(
            chemsys=query,
            fields=["material_id", "formula_pretty"]
        )
    options = [{'label': f"{res.formula_pretty} ({res.material_id})", 'value': res.material_id} for res in results]
    return options


# Callback to update the material dropdown based on search
@app.callback(
    [Output('material-dropdown', 'options'),
     Output('material-dropdown', 'value')],
    Input('search-button', 'n_clicks'),
    State('query-input', 'value'),
)
def update_material_dropdown(n_clicks, query):
    if n_clicks is None or not query:
        return [], None
    options = search_materials(query)
    if not options:
        return [], None
    return options, options[0]['value']


# Callback to display the selected material
@app.callback(
    [Output('structure-container', 'children'),
     Output('properties-container', 'children')],
    Input('display-button', 'n_clicks'),
    State('material-dropdown', 'value')
)
def display_material(n_clicks, material_id):
    if n_clicks is None or not material_id:
        return '', ''
    with MPRester(MATERIALS_PROJECT_API_KEY) as mpr:
        material = mpr.get_structure_by_material_id(material_id)
        summary = mpr.summary.get_data_by_id(material_id)

    # Create the StructureMoleculeComponent
    structure_component = ctc.StructureMoleculeComponent(material)

    # Extract key properties
    properties = {
        "Material ID": material_id,
        "Formula": summary.formula_pretty,
        "Energy Above Hull (eV/atom)": summary.energy_above_hull,
        "Space Group": summary.symmetry.symbol,
        "Band Gap (eV)": summary.band_gap,
        "Formation Energy (eV/atom)": summary.formation_energy_per_atom,
        "Magnetic Ordering": summary.ordering,
        "Total Magnetization (μB/f.u.)": summary.total_magnetization,
        "Is Stable": summary.is_stable,
        "Crystal System": summary.symmetry.crystal_system,
        "Density (g/cm³)": summary.density,
    }

    # Format properties as an HTML table
    properties_html = html.Table([
        html.Tbody([
            html.Tr([html.Th(key), html.Td(str(value))]) for key, value in properties.items()
        ])
    ], style={'border': '1px solid black', 'width': '100%', 'borderCollapse': 'collapse'})

    return structure_component.layout(), properties_html


# Register crystal toolkit with the app
ctc.register_crystal_toolkit(app, layout)

if __name__ == '__main__':
    app.run_server(debug=True, port=7860, host="0.0.0.0")