"""
MatGL Tools Module

This module provides tools for material structure relaxation and property prediction
using MatGL (Materials Graph Library) models.
"""
from __future__ import annotations
from ...core.config import material_config


import warnings
import json
from typing import Dict, List, Union, Optional, Any

import torch
from pymatgen.core import Lattice, Structure
from pymatgen.ext.matproj import MPRester
from pymatgen.io.ase import AseAtomsAdaptor

import matgl
from matgl.ext.ase import Relaxer, MolecularDynamics, PESCalculator
from ase.md.velocitydistribution import MaxwellBoltzmannDistribution

from ...core.llm_tools import llm_tool
import os
from ..support.utils import read_structure_from_file_name_or_content_string
# To suppress warnings for clearer output
warnings.simplefilter("ignore")


@llm_tool(name="relax_crystal_structure_M3GNet",
          description="Optimize crystal structure geometry using M3GNet universal potential from a structure file or content string")
async def relax_crystal_structure_M3GNet(
    structure_source: str,
    fmax: float = 0.01
) -> str:
    """
    Optimize crystal structure geometry to find its equilibrium configuration.

    Uses M3GNet universal potential for fast and accurate structure relaxation without DFT.
    Accepts a structure file or content string.

    Args:
        structure_source: The name of the structure file (e.g., POSCAR, CIF) or the content string.
        fmax: Maximum force tolerance for convergence in eV/Å (default: 0.01).

    Returns:
        A Markdown formatted string with the relaxation results or an error message.
    """
    try:
        # 使用read_structure_from_file_name_or_content_string函数读取结构
        structure_content, content_format = read_structure_from_file_name_or_content_string(structure_source)

        # 使用pymatgen读取结构
        structure = Structure.from_str(structure_content, fmt=content_format)

        if structure is None:
            return "Error: Failed to obtain a valid structure"

        # Load the M3GNet universal potential model
        pot = matgl.load_model("M3GNet-MP-2021.2.8-PES")

        # Create a relaxer and relax the structure
        relaxer = Relaxer(potential=pot)
        relax_results = relaxer.relax(structure, fmax=fmax)

        # Get the relaxed structure
        relaxed_structure = relax_results["final_structure"]
        reduced_formula = relaxed_structure.composition.reduced_formula

        # 添加结构信息
        lattice_info = relaxed_structure.lattice
        volume = relaxed_structure.volume
        density = relaxed_structure.density
        symmetry = relaxed_structure.get_space_group_info()

        # 构建原子位置表格
        sites_table = "  #  SP           a         b         c\n"
        sites_table += "---  ----  --------  --------  --------\n"
        for i, site in enumerate(relaxed_structure):
            frac_coords = site.frac_coords
            sites_table += f"{i:3d}  {site.species_string:4s}  {frac_coords[0]:8.6f}  {frac_coords[1]:8.6f}  {frac_coords[2]:8.6f}\n"

        return (f"## Structure Relaxation\n\n"
                f"- **Structure**: `{reduced_formula}`\n"
                f"- **Force Tolerance**: `{fmax} eV/Å`\n"
                f"- **Status**: `Successfully relaxed`\n\n"
                f"### Relaxed Structure Information\n\n"
                f"- **Space Group**: `{symmetry[0]} (#{symmetry[1]})`\n"
                f"- **Volume**: `{volume:.2f} ų`\n"
                f"- **Density**: `{density:.2f} g/cm³`\n"
                f"- **Lattice Parameters**:\n"
                f"  - a = `{lattice_info.a:.6f} Å`, b = `{lattice_info.b:.6f} Å`, c = `{lattice_info.c:.6f} Å`\n"
                f"  - α = `{lattice_info.alpha:.6f}°`, β = `{lattice_info.beta:.6f}°`, γ = `{lattice_info.gamma:.6f}°`\n\n"
                f"### Atomic Positions (Fractional Coordinates)\n\n"
                f"```\n"
                f"abc   :   {lattice_info.a:.6f}   {lattice_info.b:.6f}   {lattice_info.c:.6f}\n"
                f"angles:  {lattice_info.alpha:.6f}  {lattice_info.beta:.6f}  {lattice_info.gamma:.6f}\n"
                f"pbc   :       {relaxed_structure.lattice.pbc[0]!s:5s}       {relaxed_structure.lattice.pbc[1]!s:5s}       {relaxed_structure.lattice.pbc[2]!s:5s}\n"
                f"Sites ({len(relaxed_structure)})\n"
                f"{sites_table}```\n")
    except Exception as e:
        return f"Error during structure relaxation: {str(e)}"


# 内部函数，用于结构优化，返回结构对象而不是格式化字符串
async def _relax_crystal_structure_M3GNet_internal(
    structure_source: str,
    fmax: float = 0.01
) -> Union[Structure, str]:
    """
    内部使用的结构优化函数，返回结构对象而不是格式化字符串。

    Args:
        structure_source: 结构文件名或内容字符串
        fmax: 力收敛阈值 (eV/Å)

    Returns:
        优化后的结构对象或错误信息
    """
    try:
        # 使用read_structure_from_file_name_or_content_string函数读取结构
        structure_content, content_format = read_structure_from_file_name_or_content_string(structure_source)

        # 使用pymatgen读取结构
        structure = Structure.from_str(structure_content, fmt=content_format)

        if structure is None:
            return "Error: Failed to obtain a valid structure"

        # Load the M3GNet universal potential model
        pot = matgl.load_model("M3GNet-MP-2021.2.8-PES")

        # Create a relaxer and relax the structure
        relaxer = Relaxer(potential=pot)
        relax_results = relaxer.relax(structure, fmax=fmax)

        # Get the relaxed structure
        relaxed_structure = relax_results["final_structure"]

        return relaxed_structure
    except Exception as e:
        return f"Error during structure relaxation: {str(e)}"


@llm_tool(name="predict_formation_energy_M3GNet",
          description="Predict the formation energy of a crystal structure using the M3GNet formation energy model from a structure file or content string, with optional structure optimization")
async def predict_formation_energy_M3GNet(
    structure_source: str,
    optimize_structure: bool = True,
    fmax: float = 0.01
) -> str:
    """
    Predict the formation energy of a crystal structure using the M3GNet formation energy model.

    Args:
        structure_source: The name of the structure file (e.g., POSCAR, CIF) or the content string.
        optimize_structure: Whether to optimize the structure before prediction (default: True).
        fmax: Maximum force tolerance for structure relaxation in eV/Å (default: 0.01).

    Returns:
        A Markdown formatted string containing the predicted formation energy in eV/atom or an error message.
    """
    try:
        # 获取结构（优化或不优化）
        if optimize_structure:
            # 使用内部函数优化结构
            structure = await _relax_crystal_structure_M3GNet_internal(
                structure_source=structure_source,
                fmax=fmax
            )

            # 检查优化是否成功
            if isinstance(structure, str) and structure.startswith("Error"):
                return structure
        else:
            # 直接读取结构，不进行优化
            structure_content, content_format = read_structure_from_file_name_or_content_string(structure_source)
            structure = Structure.from_str(structure_content, fmt=content_format)

        if structure is None:
            return "Error: Failed to obtain a valid structure"

        # 加载预训练模型
        model = matgl.load_model("M3GNet-MP-2018.6.1-Eform")

        # 预测形成能
        eform = model.predict_structure(structure)
        reduced_formula = structure.composition.reduced_formula

        # 构建结果字符串
        optimization_status = "optimized" if optimize_structure else "non-optimized"

        # 添加结构信息
        lattice_info = structure.lattice
        volume = structure.volume
        density = structure.density
        symmetry = structure.get_space_group_info()

        # 构建原子位置表格
        sites_table = "  #  SP           a         b         c\n"
        sites_table += "---  ----  --------  --------  --------\n"
        for i, site in enumerate(structure):
            frac_coords = site.frac_coords
            sites_table += f"{i:3d}  {site.species_string:4s}  {frac_coords[0]:8.6f}  {frac_coords[1]:8.6f}  {frac_coords[2]:8.6f}\n"

        return (f"## Formation Energy Prediction\n\n"
                f"- **Structure**: `{reduced_formula}`\n"
                f"- **Structure Status**: `{optimization_status}`\n"
                f"- **Formation Energy**: `{float(eform):.3f} eV/atom`\n\n"
                f"### Structure Information\n\n"
                f"- **Space Group**: `{symmetry[0]} (#{symmetry[1]})`\n"
                f"- **Volume**: `{volume:.2f} ų`\n"
                f"- **Density**: `{density:.2f} g/cm³`\n"
                f"- **Lattice Parameters**:\n"
                f"  - a = `{lattice_info.a:.6f} Å`, b = `{lattice_info.b:.6f} Å`, c = `{lattice_info.c:.6f} Å`\n"
                f"  - α = `{lattice_info.alpha:.6f}°`, β = `{lattice_info.beta:.6f}°`, γ = `{lattice_info.gamma:.6f}°`\n\n"
                f"### Atomic Positions (Fractional Coordinates)\n\n"
                f"```\n"
                f"abc   :   {lattice_info.a:.6f}   {lattice_info.b:.6f}   {lattice_info.c:.6f}\n"
                f"angles:  {lattice_info.alpha:.6f}  {lattice_info.beta:.6f}  {lattice_info.gamma:.6f}\n"
                f"pbc   :       {structure.lattice.pbc[0]!s:5s}       {structure.lattice.pbc[1]!s:5s}       {structure.lattice.pbc[2]!s:5s}\n"
                f"Sites ({len(structure)})\n"
                f"{sites_table}```\n")
    except Exception as e:
        return f"Error: {str(e)}"


@llm_tool(name="run_molecular_dynamics_M3GNet",
          description="Run molecular dynamics simulation on a crystal structure using M3GNet universal potential, with optional structure optimization")
async def run_molecular_dynamics_M3GNet(
    structure_source: str,
    temperature_K: float = 300,
    steps: int = 100,
    optimize_structure: bool = True,
    fmax: float = 0.01
) -> str:
    """
    Run molecular dynamics simulation on a crystal structure using M3GNet universal potential.

    Args:
        structure_source: The name of the structure file (e.g., POSCAR, CIF) or the content string.
        temperature_K: Temperature for MD simulation in Kelvin (default: 300).
        steps: Number of MD steps to run (default: 100).
        optimize_structure: Whether to optimize the structure before simulation (default: True).
        fmax: Maximum force tolerance for structure relaxation in eV/Å (default: 0.01).

    Returns:
        A Markdown formatted string containing the simulation results, including final potential energy.
    """
    try:
        # 获取结构（优化或不优化）
        if optimize_structure:
            # 使用内部函数优化结构
            structure = await _relax_crystal_structure_M3GNet_internal(
                structure_source=structure_source,
                fmax=fmax
            )

            # 检查优化是否成功
            if isinstance(structure, str) and structure.startswith("Error"):
                return structure
        else:
            # 直接读取结构，不进行优化
            structure_content, content_format = read_structure_from_file_name_or_content_string(structure_source)
            structure = Structure.from_str(structure_content, fmt=content_format)

        if structure is None:
            return "Error: Failed to obtain a valid structure"

        # Load the M3GNet universal potential model
        pot = matgl.load_model("M3GNet-MP-2021.2.8-PES")

        # Convert pymatgen structure to ASE atoms
        ase_adaptor = AseAtomsAdaptor()
        atoms = ase_adaptor.get_atoms(structure)

        # Initialize the velocity according to Maxwell Boltzmann distribution
        MaxwellBoltzmannDistribution(atoms, temperature_K=temperature_K)

        # Create the MD class and run simulation
        driver = MolecularDynamics(atoms, potential=pot, temperature=temperature_K)
        driver.run(steps)

        # Get final potential energy
        final_energy = atoms.get_potential_energy()

        # Get final structure
        final_structure = ase_adaptor.get_structure(atoms)
        reduced_formula = final_structure.composition.reduced_formula

        # 构建结果字符串
        optimization_status = "optimized" if optimize_structure else "non-optimized"

        # 添加结构信息
        lattice_info = final_structure.lattice
        volume = final_structure.volume
        density = final_structure.density
        symmetry = final_structure.get_space_group_info()

        # 构建原子位置表格
        sites_table = "  #  SP           a         b         c\n"
        sites_table += "---  ----  --------  --------  --------\n"
        for i, site in enumerate(final_structure):
            frac_coords = site.frac_coords
            sites_table += f"{i:3d}  {site.species_string:4s}  {frac_coords[0]:8.6f}  {frac_coords[1]:8.6f}  {frac_coords[2]:8.6f}\n"

        return (f"## Molecular Dynamics Simulation\n\n"
                f"- **Structure**: `{reduced_formula}`\n"
                f"- **Structure Status**: `{optimization_status}`\n"
                f"- **Temperature**: `{temperature_K} K`\n"
                f"- **Steps**: `{steps}`\n"
                f"- **Final Potential Energy**: `{float(final_energy):.3f} eV`\n\n"
                f"### Final Structure Information\n\n"
                f"- **Space Group**: `{symmetry[0]} (#{symmetry[1]})`\n"
                f"- **Volume**: `{volume:.2f} ų`\n"
                f"- **Density**: `{density:.2f} g/cm³`\n"
                f"- **Lattice Parameters**:\n"
                f"  - a = `{lattice_info.a:.6f} Å`, b = `{lattice_info.b:.6f} Å`, c = `{lattice_info.c:.6f} Å`\n"
                f"  - α = `{lattice_info.alpha:.6f}°`, β = `{lattice_info.beta:.6f}°`, γ = `{lattice_info.gamma:.6f}°`\n\n"
                f"### Atomic Positions (Fractional Coordinates)\n\n"
                f"```\n"
                f"abc   :   {lattice_info.a:.6f}   {lattice_info.b:.6f}   {lattice_info.c:.6f}\n"
                f"angles:  {lattice_info.alpha:.6f}  {lattice_info.beta:.6f}  {lattice_info.gamma:.6f}\n"
                f"pbc   :       {final_structure.lattice.pbc[0]!s:5s}       {final_structure.lattice.pbc[1]!s:5s}       {final_structure.lattice.pbc[2]!s:5s}\n"
                f"Sites ({len(final_structure)})\n"
                f"{sites_table}```\n")
    except Exception as e:
        return f"Error: {str(e)}"


@llm_tool(name="calculate_single_point_energy_M3GNet",
          description="Calculate single point energy of a crystal structure using M3GNet universal potential, with optional structure optimization")
async def calculate_single_point_energy_M3GNet(
    structure_source: str,
    optimize_structure: bool = True,
    fmax: float = 0.01
) -> str:
    """
    Calculate single point energy of a crystal structure using M3GNet universal potential.

    Args:
        structure_source: The name of the structure file (e.g., POSCAR, CIF) or the content string.
        optimize_structure: Whether to optimize the structure before calculation (default: True).
        fmax: Maximum force tolerance for structure relaxation in eV/Å (default: 0.01).

    Returns:
        A Markdown formatted string containing the calculated potential energy in eV.
    """
    try:
        # 获取结构（优化或不优化）
        if optimize_structure:
            # 使用内部函数优化结构
            structure = await _relax_crystal_structure_M3GNet_internal(
                structure_source=structure_source,
                fmax=fmax
            )

            # 检查优化是否成功
            if isinstance(structure, str) and structure.startswith("Error"):
                return structure
        else:
            # 直接读取结构，不进行优化
            structure_content, content_format = read_structure_from_file_name_or_content_string(structure_source)
            structure = Structure.from_str(structure_content, fmt=content_format)

        if structure is None:
            return "Error: Failed to obtain a valid structure"

        # Load the M3GNet universal potential model
        pot = matgl.load_model("M3GNet-MP-2021.2.8-PES")

        # Convert pymatgen structure to ASE atoms
        ase_adaptor = AseAtomsAdaptor()
        atoms = ase_adaptor.get_atoms(structure)

        # Set up the calculator for atoms object
        calc = PESCalculator(pot)
        atoms.set_calculator(calc)

        # Calculate potential energy
        energy = atoms.get_potential_energy()
        reduced_formula = structure.composition.reduced_formula

        # 构建结果字符串
        optimization_status = "optimized" if optimize_structure else "non-optimized"

        # 添加结构信息
        lattice_info = structure.lattice
        volume = structure.volume
        density = structure.density
        symmetry = structure.get_space_group_info()

        # 构建原子位置表格
        sites_table = "  #  SP           a         b         c\n"
        sites_table += "---  ----  --------  --------  --------\n"
        for i, site in enumerate(structure):
            frac_coords = site.frac_coords
            sites_table += f"{i:3d}  {site.species_string:4s}  {frac_coords[0]:8.6f}  {frac_coords[1]:8.6f}  {frac_coords[2]:8.6f}\n"

        return (f"## Single Point Energy Calculation\n\n"
                f"- **Structure**: `{reduced_formula}`\n"
                f"- **Structure Status**: `{optimization_status}`\n"
                f"- **Potential Energy**: `{float(energy):.3f} eV`\n\n"
                f"### Structure Information\n\n"
                f"- **Space Group**: `{symmetry[0]} (#{symmetry[1]})`\n"
                f"- **Volume**: `{volume:.2f} ų`\n"
                f"- **Density**: `{density:.2f} g/cm³`\n"
                f"- **Lattice Parameters**:\n"
                f"  - a = `{lattice_info.a:.6f} Å`, b = `{lattice_info.b:.6f} Å`, c = `{lattice_info.c:.6f} Å`\n"
                f"  - α = `{lattice_info.alpha:.6f}°`, β = `{lattice_info.beta:.6f}°`, γ = `{lattice_info.gamma:.6f}°`\n\n"
                f"### Atomic Positions (Fractional Coordinates)\n\n"
                f"```\n"
                f"abc   :   {lattice_info.a:.6f}   {lattice_info.b:.6f}   {lattice_info.c:.6f}\n"
                f"angles:  {lattice_info.alpha:.6f}  {lattice_info.beta:.6f}  {lattice_info.gamma:.6f}\n"
                f"pbc   :       {structure.lattice.pbc[0]!s:5s}       {structure.lattice.pbc[1]!s:5s}       {structure.lattice.pbc[2]!s:5s}\n"
                f"Sites ({len(structure)})\n"
                f"{sites_table}```\n")
    except Exception as e:
        return f"Error: {str(e)}"

#Error: Bad serialized model or bad model name. It is possible that you have an older model cached. Please clear your cache by running `python -c "import matgl; matgl.clear_cache()"`
# @llm_tool(name="predict_band_gap",
#           description="Predict the band gap of a crystal structure using MEGNet multi-fidelity model from either a chemical formula or CIF file, with structure optimization")
# async def predict_band_gap(
#     formula: str = None,
#     cif_file_name: str = None,
#     method: str = "PBE",
#     fmax: float = 0.01
# ) -> str:
#     """
#     Predict the band gap of a crystal structure using the MEGNet multi-fidelity band gap model.

#     First optimizes the crystal structure using M3GNet universal potential, then predicts
#     the band gap based on the relaxed structure for more accurate results.

#     Accepts either a chemical formula (searches Materials Project database) or a CIF file.

#     Args:
#         formula: Chemical formula to retrieve from Materials Project (e.g., "Fe2O3").
#         cif_file_name: Name of CIF file in temp directory to use as structure source.
#         method: The DFT method to use for the prediction. Options are "PBE", "GLLB-SC", "HSE", or "SCAN".
#                 Default is "PBE".
#         fmax: Maximum force tolerance for structure relaxation in eV/Å (default: 0.01).

#     Returns:
#         A string containing the predicted band gap in eV or an error message.
#     """
#     try:
#         # First, relax the crystal structure
#         relaxed_result = await relax_crystal_structure(
#             formula=formula,
#             cif_file_name=cif_file_name,
#             fmax=fmax
#         )

#         # Check if relaxation was successful
#         if isinstance(relaxed_result, str) and relaxed_result.startswith("Error"):
#             return relaxed_result

#         # Use the relaxed structure for band gap prediction
#         structure = relaxed_result

#         if structure is None:
#             return "Error: Failed to obtain a valid relaxed structure"

#         # Load the pre-trained MEGNet band gap model
#         model = matgl.load_model("MEGNet-MP-2019.4.1-BandGap-mfi")

#         # Map method name to index
#         method_map = {"PBE": 0, "GLLB-SC": 1, "HSE": 2, "SCAN": 3}
#         if method not in method_map:
#             return f"Error: Unsupported method: {method}. Choose from PBE, GLLB-SC, HSE, or SCAN."

#         # Set the graph label based on the method
#         graph_attrs = torch.tensor([method_map[method]])

#         # Predict the band gap using the relaxed structure
#         bandgap = model.predict_structure(structure=structure, state_attr=graph_attrs)
#         reduced_formula = structure.reduced_formula

#         # Return the band gap as a string
#         return f"The predicted band gap for relaxed {reduced_formula} using {method} method is {float(bandgap):.3f} eV."
#     except Exception as e:
#         return f"Error: {str(e)}"