"""
Property Prediction Module

This module provides functions for predicting properties of crystal structures.
"""

import asyncio
import torch
import numpy as np
from ase.units import GPa
from mattersim.forcefield import MatterSimCalculator

from mars_toolkit.core.llm_tools import llm_tool
from mars_toolkit.compute.structure_opt import convert_structure

@llm_tool(
    name="predict_properties",
    description="Predict energy, forces, and stress of crystal structures based on CIF string",
)
async def predict_properties(cif_content: str) -> str:
    """
    Use MatterSim to predict energy, forces, and stress of crystal structures.

    Args:
        cif_content: Crystal structure string in CIF format

    Returns:
        String containing prediction results
    """
    # 使用asyncio.to_thread异步执行可能阻塞的操作
    def run_prediction():
        # 使用 convert_structure 函数将 CIF 字符串转换为 Atoms 对象
        structure = convert_structure("cif", cif_content)
        if structure is None:
            return "Unable to parse CIF string. Please check if the format is correct."

        # 设置设备
        device = "cuda" if torch.cuda.is_available() else "cpu"

        # 使用 MatterSimCalculator 计算属性
        structure.calc = MatterSimCalculator(device=device)

        # 直接获取能量、力和应力
        energy = structure.get_potential_energy()
        forces = structure.get_forces()
        stresses = structure.get_stress(voigt=False)

        # 计算每原子能量
        num_atoms = len(structure)
        energy_per_atom = energy / num_atoms

        # 计算应力（GPa和eV/A^3格式）
        stresses_ev_a3 = stresses
        stresses_gpa = stresses / GPa

        # 构建返回的提示信息
        prompt = f"""
## {structure.get_chemical_formula()} Crystal Structure Property Prediction Results

Prediction results using the provided CIF structure:

- Total Energy (eV): {energy}
- Energy per Atom (eV/atom): {energy_per_atom:.4f}
- Forces (eV/Angstrom): {forces[0]}  # Forces on the first atom
- Stress (GPa): {stresses_gpa[0][0]}  # First component of the stress tensor
- Stress (eV/A^3): {stresses_ev_a3[0][0]}  # First component of the stress tensor

"""
        return prompt

    # 异步执行预测操作
    return await asyncio.to_thread(run_prediction)