# !/usr/bin/env python3

import os, warnings, random, shutil, re
import numpy as np
from pathlib import Path
from dotenv import load_dotenv
from ase.io import read, write
from ase.build import surface
from icet import ClusterSpace
from icet.tools.structure_generation import generate_sqs
from icet.input_output.logging_tools import set_log_config
from mp_api.client import MPRester
from pymatgen.core import Structure
from pymatgen.io.vasp import Poscar, Kpoints
from pymatgen.analysis.defects import generators
from pymatgen.io.vasp.sets import (
    MPStaticSet, 
    MPRelaxSet, 
    MPNonSCFSet, 
    MPScanRelaxSet, 
    MPScanStaticSet,
    MPMDSet,
    NEBSet,
    MVLElasticSet,
    )

from masgent import schemas
from masgent.utils.interface_maker import run_interface_maker
from masgent.utils.utils import (
    write_comments,
    color_print,
    ask_for_mp_api_key,
    validate_mp_api_key,
    generate_batch_script,
    list_files_in_dir,
    )

# Do not show warnings
warnings.filterwarnings('ignore')

# Track whether Materials Project key has been checked during this process
_mp_key_checked = False

def with_metadata(input: schemas.ToolMetadata):
    '''
    Decorator to add metadata to tool functions.
    '''
    def decorator(func):
        func._tool_metadata = input
        return func
    return decorator

@with_metadata(schemas.ToolMetadata(
    name='List Files',
    description='List all files in the current session runs directory.',
    requires=[],
    optional=[],
    defaults={},
    prereqs=[],
))
def list_files() -> dict:
    runs_dir = os.environ.get('MASGENT_SESSION_RUNS_DIR')

    file_list = []
    base_dir = Path(runs_dir)
    for item in base_dir.rglob('*'):
        if item.is_file():
            file_list.append(str(item.relative_to(base_dir)))
    return {
        'status': 'success',
        'message': f'Found {len(file_list)} files in the current session runs directory.',
        'files': file_list,
    }

@with_metadata(schemas.ToolMetadata(
    name='Read File',
    description='Read a file from the current session runs directory.',
    requires=['name'],
    optional=[],
    defaults={},
    prereqs=[],
))
def read_file(name: str) -> dict:
    runs_dir = os.environ.get('MASGENT_SESSION_RUNS_DIR')
    base_dir = Path(runs_dir)

    try:
        with open(base_dir / name, "r") as f:
            content = f.read()
        return {
            'status': 'success',
            'message': f'File {name} read successfully.',
            'content': content,
        }
    except Exception as e:
        return {
            'status': 'error',
            'message': f'An error occurred while reading file {name}: {e}',
        }

@with_metadata(schemas.ToolMetadata(
    name='Rename File',
    description='Rename a file in the current session runs directory.',
    requires=['name', 'new_name'],
    optional=[],
    defaults={},
    prereqs=[],
))
def rename_file(name: str, new_name: str) -> dict:
    runs_dir = os.environ.get('MASGENT_SESSION_RUNS_DIR')
    base_dir = Path(runs_dir)

    try:
        new_path = base_dir / new_name
        if not str(new_path).startswith(str(base_dir)):
            return {
                'status': 'error',
                'message': f'Renaming to {new_name} would move the file outside the session runs directory, which is not allowed.',
            }

        os.makedirs(new_path.parent, exist_ok=True)
        shutil.copy2(base_dir / name, new_path)
        return {
            'status': 'success',
            'message': f'File {name} renamed to {new_name} successfully.',
        }
    except Exception as e:
        return {
            'status': 'error',
            'message': f'An error occurred while renaming file {name} to {new_name}: {e}',
        }

@with_metadata(schemas.ToolMetadata(
    name='Generate POSCARs from Materials Project',
    description='Generate all possible POSCAR files from Materials Project database based on chemical formula, with the most stable structure saved as POSCAR by default and all matched structures saved in a separate folder.',
    requires=['formula'],
    optional=[],
    defaults={},
    prereqs=[],
))
def generate_vasp_poscar(input: schemas.GenerateVaspPoscarSchema) -> dict:
    '''
    Generate VASP POSCAR file from Materials Project database.
    '''
    # color_print(f'\n[Debug: Function Calling] generate_vasp_poscar with input: {input}', 'green')
    
    formula = input.formula

    try:
        runs_dir = os.environ.get('MASGENT_SESSION_RUNS_DIR')

        poscars_dir = os.path.join(runs_dir, f'POSCARs')
        os.makedirs(poscars_dir, exist_ok=True)

        # Ensure Materials Project API key exists and validate it only once per process
        load_dotenv(dotenv_path='.env')

        global _mp_key_checked
        if not _mp_key_checked:
            if 'MP_API_KEY' not in os.environ:
                ask_for_mp_api_key()
            else:
                # color_print('[Info] Materials Project API key found in environment.\n', 'green')
                validate_mp_api_key(os.environ['MP_API_KEY'])
            _mp_key_checked = True
        
        with MPRester(mute_progress_bars=True) as mpr:
            docs = mpr.materials.summary.search(formula=formula)
            if not docs:
                return {
                    'status': 'error',
                    'message': f'No materials found in Materials Project database for formula: {formula}'
                }
        
        # Save the most stable structure in the runs directory
        mid_0 = docs[0].material_id
        structure_0 = mpr.get_structure_by_material_id(mid_0)
        poscar_0 = Poscar(structure_0)
        poscar_0.write_file(os.path.join(runs_dir, 'POSCAR'), direct=True)
        comments_0 = f'# (Most Stable) Generated by Masgent from Materials Project entry {mid_0}, crystal system: {docs[0].symmetry.crystal_system}, space group: {docs[0].symmetry.symbol}.'
        write_comments(os.path.join(runs_dir, 'POSCAR'), 'poscar', comments_0)
        
        # Save all matched structures in the poscars directory
        for doc in docs:
            mid = doc.material_id
            crystal_system = doc.symmetry.crystal_system
            space_group_symbol = doc.symmetry.symbol
            structure = mpr.get_structure_by_material_id(mid)
            poscar = Poscar(structure)

            # If "/" in space group symbol, replace with "_"
            space_group_symbol_ = space_group_symbol.replace('/', '_')
            poscar.write_file(os.path.join(poscars_dir, f'POSCAR_{crystal_system}_{space_group_symbol_}_{mid}'), direct=True)

            comments = f'# Generated by Masgent from Materials Project entry {mid}, crystal system: {crystal_system}, space group: {space_group_symbol}.'
            write_comments(os.path.join(poscars_dir, f'POSCAR_{crystal_system}_{space_group_symbol_}_{mid}'), 'poscar', comments)
        
        poscar_files = list_files_in_dir(poscars_dir) + [os.path.join(runs_dir, 'POSCAR')]
        
        return {
            'status': 'success',
            'message': f'Generated POSCAR(s) in {poscars_dir}.',
            'poscar_dir': poscars_dir,
            'all_poscars': poscar_files,
            'most_stable_poscar': os.path.join(runs_dir, 'POSCAR'),
        }

    except Exception as e:
        return {
            'status': 'error',
            'message': f'POSCAR generation failed: {str(e)}'
        }

@with_metadata(schemas.ToolMetadata(
    name='Generate VASP Inputs (INCAR, KPOINTS, POTCAR, POSCAR)',
    description='Generate VASP input files (INCAR, KPOINTS, POTCAR, POSCAR) from a given POSCAR file using pymatgen input sets (MPRelaxSet, MPStaticSet, MPNonSCFSet, MPScanRelaxSet, MPScanStaticSet, MPMDSet, NEBSet, MVLElasticSet).',
    requires=['vasp_input_sets'],
    optional=['poscar_path', 'only_incar'],
    defaults={
        'poscar_path': f'{os.environ.get("MASGENT_SESSION_RUNS_DIR")}/POSCAR',
        'only_incar': False,
        },
    prereqs=[],
))
def generate_vasp_inputs_from_poscar(input: schemas.GenerateVaspInputsFromPoscar) -> dict:
    '''
    Generate VASP input files (INCAR, KPOINTS, POTCAR, POSCAR) from a given POSCAR file using pymatgen input sets.
    '''
    # color_print(f'\n[Debug: Function Calling] generate_vasp_inputs_from_poscar with input: {input}', 'green')
    
    poscar_path = input.poscar_path
    vasp_input_sets = input.vasp_input_sets
    only_incar = input.only_incar

    VIS_MAP = {
        'MPRelaxSet': MPRelaxSet,
        'MPStaticSet': MPStaticSet,
        'MPNonSCFSet': MPNonSCFSet,
        'MPScanRelaxSet': MPScanRelaxSet,
        'MPScanStaticSet': MPScanStaticSet,
        'MPMDSet': MPMDSet,
        'NEBSet': NEBSet,
        'MVLElasticSet': MVLElasticSet,
    }
    vis_class = VIS_MAP[vasp_input_sets]

    try:
        runs_dir = os.environ.get('MASGENT_SESSION_RUNS_DIR')

        vasp_inputs_dir = os.path.join(runs_dir, f'vasp_inputs/{vasp_input_sets}')
        os.makedirs(vasp_inputs_dir, exist_ok=True)

        structure = Structure.from_file(poscar_path)
        vis = vis_class(structure)

        if only_incar:
            vis.incar.write_file(os.path.join(vasp_inputs_dir, 'INCAR'))
            vis.poscar.write_file(os.path.join(vasp_inputs_dir, 'POSCAR'))
            incar_comments = f'# Generated by Masgent using {vasp_input_sets} set provided by Materials Project.'
            write_comments(os.path.join(vasp_inputs_dir, 'INCAR'), 'incar', incar_comments)
            return {
                'status': 'success',
                'message': f'Generated INCAR based on {vasp_input_sets} in {vasp_inputs_dir}.',
                'incar_path': os.path.join(vasp_inputs_dir, 'INCAR'),
            }
        
        vis.incar.write_file(os.path.join(vasp_inputs_dir, 'INCAR'))
        vis.poscar.write_file(os.path.join(vasp_inputs_dir, 'POSCAR'))
        vis.kpoints.write_file(os.path.join(vasp_inputs_dir, 'KPOINTS'))
        vis.potcar.write_file(os.path.join(vasp_inputs_dir, 'POTCAR'))

        incar_comments = f'# Generated by Masgent using {vasp_input_sets} set provided by Materials Project.'
        write_comments(os.path.join(vasp_inputs_dir, 'INCAR'), 'incar', incar_comments)
        poscar_comments = f'# Generated by Masgent using {vasp_input_sets} set provided by Materials Project.'
        write_comments(os.path.join(vasp_inputs_dir, 'POSCAR'), 'poscar', poscar_comments)
        kpoints_comments = f'# Generated by Masgent using {vasp_input_sets} set provided by Materials Project.'
        write_comments(os.path.join(vasp_inputs_dir, 'KPOINTS'), 'kpoints', kpoints_comments)
        
        return {
            'status': 'success',
            'message': f'Generated VASP input files based on {vasp_input_sets} in {vasp_inputs_dir}.',
            'incar_path': os.path.join(vasp_inputs_dir, 'INCAR'),
            'kpoints_path': os.path.join(vasp_inputs_dir, 'KPOINTS'),
            'potcar_path': os.path.join(vasp_inputs_dir, 'POTCAR'),
            'poscar_path': os.path.join(vasp_inputs_dir, 'POSCAR'),
        }
    
    except Exception as e:
        return {
            'status': 'error',
            'message': f'VASP input files generation failed: {str(e)}'
        }

@with_metadata(schemas.ToolMetadata(
    name='Generate HPC Slurm Script',
    description='Generate HPC Slurm job submission script for VASP calculations.',
    requires=[],
    optional=['partition', 'nodes', 'ntasks', 'walltime', 'jobname', 'mail_type', 'mail_user', 'command'],
    defaults={
        'partition': 'normal',
        'nodes': 1,
        'ntasks': 8,
        'walltime': '01:00:00',
        'jobname': 'masgent_job',
        'command': 'srun vasp_std > vasp.out'
        },
    prereqs=[],
    ))
def generate_vasp_inputs_hpc_slurm_script(input: schemas.GenerateVaspInputsHpcSlurmScript) -> dict:
    '''
    Generate HPC Slurm job submission script for VASP calculations.
    '''
    # color_print(f'\n[Debug: Function Calling] generate_vasp_inputs_hpc_slurm_script with input: {input}', 'green')
    
    partition = input.partition
    nodes = input.nodes
    ntasks = input.ntasks
    walltime = input.walltime
    jobname = input.jobname
    command = input.command

    try:
        runs_dir = os.environ.get('MASGENT_SESSION_RUNS_DIR')
        
        scripts = f'''#!/bin/bash
#SBATCH --partition={partition}
#SBATCH --nodes={nodes}
#SBATCH --ntasks={ntasks}
#SBATCH --time={walltime}
#SBATCH --job-name={jobname}
#SBATCH --output={jobname}.out
#SBATCH --error={jobname}.err

# This Slurm script was generated by Masgent, customize as needed.
{command}
'''
        script_path = os.path.join(runs_dir, 'masgent_submit.sh')
        with open(script_path, 'w') as f:
            f.write(scripts)

        return {
            'status': 'success',
            'message': f'Generated Slurm script in {script_path}.',
            'script_path': script_path,
        }
    
    except Exception as e:
        return {
            'status': 'error',
            'message': f'Slurm script generation failed: {str(e)}'
        }

@with_metadata(schemas.ToolMetadata(
    name='Convert Structure Format',
    description='Convert structure files between different formats (CIF, POSCAR, XYZ).',
    requires=['input_path', 'input_format', 'output_format'],
    optional=[],
    defaults={},
    prereqs=[],
))
def convert_structure_format(input: schemas.ConvertStructureFormatSchema) -> dict:
    '''
    Convert structure files between different formats (CIF, POSCAR, XYZ).
    '''
    # color_print(f'\n[Debug: Function Calling] convert_structure_format with input: {input}', 'green')
    
    input_path = input.input_path
    input_format = input.input_format
    output_format = input.output_format
    
    format_map = {
        "POSCAR": "vasp",
        "CIF": "cif",
        "XYZ": "xyz"
    }

    try:
        runs_dir = os.environ.get('MASGENT_SESSION_RUNS_DIR')

        convert_dir = os.path.join(runs_dir, 'convert')
        os.makedirs(convert_dir, exist_ok=True)

        atoms = read(input_path, format=format_map[input_format])
        filename_wo_ext = os.path.splitext(os.path.basename(input_path))[0]
        # Ignore the POSCAR, do not add extension
        if output_format == 'POSCAR':
            output_path = os.path.join(convert_dir, 'POSCAR')
        else:
            output_path = os.path.join(convert_dir, f'{filename_wo_ext}.{output_format.lower()}')
        write(output_path, atoms, format=format_map[output_format])

        return {
            'status': 'success',
            'message': f'Converted structure saved to {output_path}.',
            'output_path': output_path
        }
    
    except Exception as e:
        return {
            'status': 'error',
            'message': f'Structure conversion failed: {str(e)}'
        }

@with_metadata(schemas.ToolMetadata(
    name='Convert POSCAR Coordinates',
    description='Convert POSCAR between direct and cartesian coordinates.',
    requires=['to_cartesian'],
    optional=['poscar_path'],
    defaults={'poscar_path': f'{os.environ.get("MASGENT_SESSION_RUNS_DIR")}/POSCAR'},
    prereqs=[],
))
def convert_poscar_coordinates(input: schemas.ConvertPoscarCoordinatesSchema) -> dict:
    '''
    Convert POSCAR between direct and cartesian coordinates.
    '''
    # color_print(f'\n[Debug: Function Calling] convert_poscar_coordinates with input: {input}', 'green')
    
    poscar_path = input.poscar_path
    to_cartesian = input.to_cartesian

    try:
        runs_dir = os.environ.get('MASGENT_SESSION_RUNS_DIR')

        convert_dir = os.path.join(runs_dir, 'convert')
        os.makedirs(convert_dir, exist_ok=True)
        
        structure = Structure.from_file(poscar_path)
        poscar = Poscar(structure)
        poscar.write_file(os.path.join(convert_dir, 'POSCAR'), direct=not to_cartesian)
        coord_type = 'Cartesian' if to_cartesian else 'Direct'
        comments = f'# Generated by Masgent converted to {coord_type} coordinates.'
        write_comments(os.path.join(convert_dir, 'POSCAR'), 'poscar', comments)

        return {
            'status': 'success',
            'message': f'Converted POSCAR to {coord_type} coordinates in {convert_dir}.',
            'poscar_path': os.path.join(convert_dir, 'POSCAR'),
        }

    except Exception as e:
        return {
            'status': 'error',
            'message': f'POSCAR coordinate conversion failed: {str(e)}'
        }

@with_metadata(schemas.ToolMetadata(
    name='Customize KPOINTS',
    description='Customize VASP KPOINTS from POSCAR with specified accuracy level (Low, Medium, High).',
    requires=['accuracy_level'],
    optional=['poscar_path'],
    defaults={'poscar_path': f'{os.environ.get("MASGENT_SESSION_RUNS_DIR")}/POSCAR'},
    prereqs=[],
))
def customize_vasp_kpoints_with_accuracy(input: schemas.CustomizeVaspKpointsWithAccuracy) -> dict:
    '''
    Customize VASP KPOINTS from POSCAR with specified accuracy level.
    '''
    # color_print(f'\n[Debug: Function Calling] customize_vasp_kpoints_with_accuracy with input: {input}', 'green')
    
    poscar_path = input.poscar_path
    accuracy_level = input.accuracy_level
    
    DENSITY_MAP = {
        'Low': 1000,
        'Medium': 3000,
        'High': 5000,
    }
    kppa = DENSITY_MAP[accuracy_level]

    try:
        runs_dir = os.environ.get('MASGENT_SESSION_RUNS_DIR')
        
        structure = Structure.from_file(poscar_path)
        kpoints = Kpoints.automatic_density(structure, kppa=kppa)
        kpoints.write_file(os.path.join(runs_dir, 'KPOINTS'))
        comments = f'# Generated by Masgent with {accuracy_level} accuracy (Grid Density = {kppa} / number of atoms)'
        write_comments(os.path.join(runs_dir, 'KPOINTS'), 'kpoints', comments)
        
        return {
            'status': 'success',
            'message': f'Updated KPOINTS with {accuracy_level} accuracy in {runs_dir}.',
            'kpoints_path': os.path.join(runs_dir, 'KPOINTS'),
        }
    
    except Exception as e:
        return {
            'status': 'error',
            'message': f'VASP KPOINTS generation failed: {str(e)}'
        }

@with_metadata(schemas.ToolMetadata(
    name='Generate Vacancy Defects',
    description='Generate VASP POSCAR with vacancy defects.',
    requires=['original_element', 'defect_amount'],
    optional=['poscar_path'],
    defaults={'poscar_path': f'{os.environ.get("MASGENT_SESSION_RUNS_DIR")}/POSCAR'},
    prereqs=[],
))
def generate_vasp_poscar_with_vacancy_defects(input: schemas.GenerateVaspPoscarWithVacancyDefects) -> dict:
    '''
    Generate VASP POSCAR with vacancy defects.
    '''
    # color_print(f'\n[Debug: Function Calling] generate_vasp_poscar_with_vacancy_defects with input: {input}', 'green')
    
    poscar_path = input.poscar_path
    original_element = input.original_element
    defect_amount = input.defect_amount

    try:
        runs_dir = os.environ.get('MASGENT_SESSION_RUNS_DIR')

        defect_dir = os.path.join(runs_dir, 'defects')
        os.makedirs(defect_dir, exist_ok=True)
        
        atoms = read(poscar_path, format='vasp')

        all_indices = [i for i, atom in enumerate(atoms) if atom.symbol == original_element]
        if isinstance(defect_amount, float):
            num_defects = max(1, int(defect_amount * len(all_indices)))
        elif isinstance(defect_amount, int):
            num_defects = defect_amount

        vacancy_indices = random.sample(all_indices, num_defects)
        del atoms[vacancy_indices]

        write(os.path.join(defect_dir, 'POSCAR'), atoms, format='vasp', direct=True, sort=True)
        comments = f'# Generated by Masgent with vacancy defects of element {original_element} by randomly removing {num_defects} atoms, be careful to verify structure.'
        write_comments(os.path.join(defect_dir, 'POSCAR'), 'poscar', comments)

        # return f'\nGenerated POSCAR with vacancy defects in {os.path.join(target_dir, "POSCAR")}.'
        return {
            'status': 'success',
            'message': f'Generated POSCAR with vacancy defects in {os.path.join(defect_dir, "POSCAR")}.',
            'poscar_path': os.path.join(defect_dir, 'POSCAR'),
        }
    
    except Exception as e:
        return {
            'status': 'error',
            'message': f'VASP POSCAR defect generation failed: {str(e)}'
        }

@with_metadata(schemas.ToolMetadata(
    name='Generate Substitution Defects',
    description='Generate VASP POSCAR with substitution defects.',
    requires=['original_element', 'defect_element', 'defect_amount'],
    optional=['poscar_path'],
    defaults={'poscar_path': f'{os.environ.get("MASGENT_SESSION_RUNS_DIR")}/POSCAR'},
    prereqs=[],
))
def generate_vasp_poscar_with_substitution_defects(input: schemas.GenerateVaspPoscarWithSubstitutionDefects) -> dict:
    '''
    Generate VASP POSCAR with substitution defects.
    '''
    # color_print(f'\n[Debug: Function Calling] generate_vasp_poscar_with_substitution_defects with input: {input}', 'green')

    poscar_path = input.poscar_path
    original_element = input.original_element
    defect_element = input.defect_element
    defect_amount = input.defect_amount

    try:
        runs_dir = os.environ.get('MASGENT_SESSION_RUNS_DIR')

        defect_dir = os.path.join(runs_dir, 'defects')
        os.makedirs(defect_dir, exist_ok=True)
        
        atoms = read(poscar_path, format='vasp')

        all_indices = [i for i, atom in enumerate(atoms) if atom.symbol == original_element]
        if isinstance(defect_amount, float):
            num_defects = max(1, int(defect_amount * len(all_indices)))
        elif isinstance(defect_amount, int):
            num_defects = defect_amount

        substitution_indices = random.sample(all_indices, num_defects)
        for i in substitution_indices:
            atoms[i].symbol = defect_element
        
        write(os.path.join(defect_dir, 'POSCAR'), atoms, format='vasp', direct=True, sort=True)
        comments = f'# Generated by Masgent with substitution defect of element {original_element} to {defect_element} by randomly substituting {num_defects} atoms, be careful to verify structure.'
        write_comments(os.path.join(defect_dir, 'POSCAR'), 'poscar', comments)

        return {
            'status': 'success',
            'message': f'Generated POSCAR with substitution defects in {os.path.join(defect_dir, "POSCAR")}.',
            'poscar_path': os.path.join(defect_dir, 'POSCAR'),
        }
    
    except Exception as e:
        return {
            'status': 'error',
            'message': f'VASP POSCAR defect generation failed: {str(e)}'
        }

@with_metadata(schemas.ToolMetadata(
    name='Generate Interstitial (Voronoi) Defects',
    description='Generate VASP POSCAR with interstitial (Voronoi) defects.',
    requires=['defect_element'],
    optional=['poscar_path'],
    defaults={'poscar_path': f'{os.environ.get("MASGENT_SESSION_RUNS_DIR")}/POSCAR'},
    prereqs=[],
))
def generate_vasp_poscar_with_interstitial_defects(input: schemas.GenerateVaspPoscarWithInterstitialDefects) -> dict:
    '''
    Generate VASP POSCAR with interstitial (Voronoi) defects.
    '''
    # color_print(f'\n[Debug: Function Calling] generate_vasp_poscar_with_interstitial_defects with input: {input}', 'green')

    poscar_path = input.poscar_path
    defect_element = input.defect_element

    try:
        runs_dir = os.environ.get('MASGENT_SESSION_RUNS_DIR')

        defect_dir = os.path.join(runs_dir, 'defects')
        os.makedirs(defect_dir, exist_ok=True)
        
        atoms = read(poscar_path, format='vasp')

        # Read atoms from ASE and convert to Pymatgen Structure
        structure = Structure.from_ase_atoms(atoms)
        interstitial_generator = generators.VoronoiInterstitialGenerator().generate(structure=structure, insert_species=[defect_element])
        defect_sites, defect_structures = [], []
        for defect in interstitial_generator:
            defect_sites.append(defect.site.frac_coords)
            defect_structures.append(defect.defect_structure)

        if len(defect_structures) == 0:
            return {
                'status': 'error',
                'message': f'No interstitial sites found for element {defect_element}.',
            }
        else:
            for i, defect_structure in enumerate(defect_structures):
                # Convert back to ASE Atoms for writing
                defect_atoms = defect_structure.to_ase_atoms()
                write(os.path.join(defect_dir, f'POSCAR_{i}'), defect_atoms, format='vasp', direct=True, sort=True)
                comments = f'# Generated by Masgent with interstitial (Voronoi) defect of element {defect_element} at fract. coords {defect_sites[i]}, be careful to verify structure.'
                write_comments(os.path.join(defect_dir, f'POSCAR_{i}'), 'poscar', comments)

        # return f'\nGenerated POSCAR with interstitial (Voronoi) defects in {target_dir}.'
        return {
            'status': 'success',
            'message': f'Generated POSCAR(s) with interstitial (Voronoi) defects in {defect_dir}.',
            'poscar_paths': [os.path.join(defect_dir, f'POSCAR_{i}') for i in range(len(defect_structures))],
        }
    
    except Exception as e:
        return {
            'status': 'error',
            'message': f'VASP POSCAR defect generation failed: {str(e)}'
        }

@with_metadata(schemas.ToolMetadata(
    name='Generate Supercell',
    description='Generate supercell from POSCAR based on user-defined 3x3 scaling matrix.',
    requires=['scaling_matrix'],
    optional=['poscar_path'],
    defaults={'poscar_path': f'{os.environ.get("MASGENT_SESSION_RUNS_DIR")}/POSCAR'},
    prereqs=[],
))
def generate_supercell_from_poscar(input: schemas.GenerateSupercellFromPoscar) -> dict:
    '''
    Generate supercell from POSCAR based on user-defined 3x3 scaling matrix.
    '''
    # color_print(f'\n[Debug: Function Calling] generate_supercell_from_poscar with input: {input}', 'green')
    
    poscar_path = input.poscar_path
    scaling_matrix = input.scaling_matrix

    scaling_matrix_ = [
        [int(num) for num in line.strip().split()] 
        for line in scaling_matrix.split(';')
        ]

    try:
        runs_dir = os.environ.get('MASGENT_SESSION_RUNS_DIR')

        supercell_dir = os.path.join(runs_dir, 'supercell')
        os.makedirs(supercell_dir, exist_ok=True)
        
        structure = Structure.from_file(poscar_path).copy()
        supercell_structure = structure.make_supercell(scaling_matrix_)
        supercell_poscar = Poscar(supercell_structure)
        supercell_poscar.write_file(os.path.join(supercell_dir, 'POSCAR'), direct=True)
        
        comments = f'# Generated by Masgent as supercell with scaling matrix {scaling_matrix}.'
        write_comments(os.path.join(supercell_dir, 'POSCAR'), 'poscar', comments)

        return {
            'status': 'success',
            'message': f'Generated supercell POSCAR in {os.path.join(supercell_dir, "POSCAR")}.',
            'poscar_path': os.path.join(supercell_dir, 'POSCAR'),
        }
    
    except Exception as e:
        return {
            'status': 'error',
            'message': f'Supercell generation failed: {str(e)}'
        }

@with_metadata(schemas.ToolMetadata(
    name='Generate Special Quasirandom Structures (SQS)',
    description='Generate Special Quasirandom Structures (SQS) using icet based on given POSCAR',
    requires=['target_sites', 'target_concentrations'],
    optional=['poscar_path', 'cutoffs', 'max_supercell_size', 'mc_steps'],
    defaults={
        'poscar_path': f'{os.environ.get("MASGENT_SESSION_RUNS_DIR")}/POSCAR',
        'cutoffs': [8.0, 4.0],
        'max_supercell_size': 8,
        'mc_steps': 10000,
    },
))
def generate_sqs_from_poscar(input: schemas.GenerateSqsFromPoscar) -> dict:
    '''
    Generate Special Quasirandom Structures (SQS) using icet.
    '''

    poscar_path = input.poscar_path # e.g., 'LaCoO3'
    target_configurations = input.target_configurations # e.g., {'La': {'La': 0.5, 'Y': 0.5}, 'Co': {'Al': 0.75, 'Co': 0.25}}
    cutoffs = input.cutoffs
    max_supercell_size = input.max_supercell_size
    mc_steps = input.mc_steps

    try:
        runs_dir = os.environ.get('MASGENT_SESSION_RUNS_DIR')

        sqs_dir = os.path.join(runs_dir, 'sqs')
        os.makedirs(sqs_dir, exist_ok=True)

        set_log_config(
        filename=f'{sqs_dir}/masgent_sqs.log',
        level='INFO',
        )

        primitive_structure = read(poscar_path, format='vasp')

        # Get the all chemical symbols in the structure
        chem_symbols = [[site] for site in primitive_structure.get_chemical_symbols()]

        # Create target sites mapping based on target_configurations: {'La': ['La', 'Y'], 'Co': ['Al', 'Co']}
        target_sites = {}
        for site, config in target_configurations.items():
            target_sites[site] = list(config.keys())

        # Update chem_symbols to reflect target sites
        for i, site in enumerate(chem_symbols):
            for target_site, configurations in target_sites.items():
                if site == [target_site]:
                    chem_symbols[i] = configurations

        # Initialize ClusterSpace
        cs = ClusterSpace(
            structure=primitive_structure, 
            cutoffs=cutoffs, 
            chemical_symbols=chem_symbols,
            )
        
        # Get the sublattice (A, B, ...) configurations from ClusterSpace
        chemical_symbol_representation = cs._get_chemical_symbol_representation()
        
        # Map sublattice letters to actual element symbols
        sublattice_indices = {}
        sublattice_parts = chemical_symbol_representation.split('),')
        for i, part in enumerate(sublattice_parts):
            match = re.search(r"\['(.*)'\]", part)
            if match:
                elements = match.group(1).split("', '")
                sublattice_indices[chr(65 + i)] = elements
        
        # Initialize target concentrations: {'A': {'Al': 0.0, 'Co': 0.0}, 'B': {'La': 0.0, 'Y': 0.0}, 'O': {'O': 1.0}}
        unique_chem_symbols = [list(x) for x in dict.fromkeys(tuple(x) for x in chem_symbols)]
        target_concentrations = {}
        for sublattice, elements in sublattice_indices.items():
            for unique_list in unique_chem_symbols:
                if set(elements) == set(unique_list):
                    concentration_dict = {element: 0.0 for element in unique_list}
                    target_concentrations[sublattice] = concentration_dict
                if len(unique_list) == 1:
                    element = unique_list[0]
                    target_concentrations[element] = {element: 1.0}

        # Update target concentrations based on target_configurations
        for key, value in target_configurations.items():
            for sublattice, conc_dict in target_concentrations.items():
                if key in conc_dict:
                    target_concentrations[sublattice] = value

        # Generate SQS and write to POSCAR
        sqs = generate_sqs(cluster_space=cs,
                   max_size=max_supercell_size,
                   target_concentrations=target_concentrations,
                   n_steps=mc_steps
                   )
        write(os.path.join(sqs_dir, 'POSCAR'), sqs, format='vasp', direct=True, sort=True)
        comments = f'# Generated by Masgent as Special Quasirandom Structure (SQS) with target configurations {target_configurations} using icet.'
        write_comments(os.path.join(sqs_dir, 'POSCAR'), 'poscar', comments)

        return {
            'status': 'success',
            'message': f'Generated SQS POSCAR in {os.path.join(sqs_dir, "POSCAR")}.',
            'poscar_path': os.path.join(sqs_dir, 'POSCAR'),
        }

    except Exception as e:
        return {
            'status': 'error',
            'message': f'SQS generation failed: {str(e)}'
        }
    
@with_metadata(schemas.ToolMetadata(
    name='Generate suface slab from bulk POSCAR',
    description='Generate surface slab from bulk POSCAR based on Miller indices, vacuum thickness, and slab layers',
    requires=['miller_indices'],
    optional=['poscar_path', 'vacuum_thickness', 'slab_layers'],
    defaults={
        'poscar_path': f'{os.environ.get("MASGENT_SESSION_RUNS_DIR")}/POSCAR',
        'vacuum_thickness': 15.0,
        'slab_layers': 4,
        },
    prereqs=[],
))
def generate_surface_slab_from_poscar(input: schemas.GenerateSurfaceSlabFromPoscar) -> dict:
    '''
    Generate VASP POSCAR for surface slab from bulk POSCAR based on Miller indices, vacuum thickness, and slab layers
    '''
    # color_print(f'\n[Debug: Function Calling] generate_vasp_poscar_for_surface_slab with input: {input}', 'green')
    
    poscar_path = input.poscar_path
    miller_indices = input.miller_indices
    vacuum_thickness = input.vacuum_thickness
    slab_layers = input.slab_layers

    try:
        runs_dir = os.environ.get('MASGENT_SESSION_RUNS_DIR')
        
        surface_slab_dir = os.path.join(runs_dir, 'surface_slab')
        os.makedirs(surface_slab_dir, exist_ok=True)
        
        bulk_atoms = read(poscar_path, format='vasp')
        slab_atoms = surface(lattice=bulk_atoms, indices=miller_indices, layers=slab_layers, vacuum=vacuum_thickness, tol=1e-10, periodic=True)
        write(os.path.join(surface_slab_dir, 'POSCAR'), slab_atoms, format='vasp', direct=True, sort=True)
        comments = f'# Generated by Masgent as surface slab with Miller indices {miller_indices}, vacuum thickness {vacuum_thickness} Å, and slab layers {slab_layers}.'
        write_comments(os.path.join(surface_slab_dir, 'POSCAR'), 'poscar', comments)

        return {
            'status': 'success',
            'message': f'Generated surface slab POSCAR in {os.path.join(surface_slab_dir, "POSCAR")}.',
            'poscar_path': os.path.join(surface_slab_dir, 'POSCAR'),
        }

    except Exception as e:
        return {
            'status': 'error',
            'message': f'Surface slab POSCAR generation failed: {str(e)}'
        }

@with_metadata(schemas.ToolMetadata(
    name='Generate interface from two POSCARs',
    description='Generate VASP POSCAR for interface from two given POSCAR files based on specified parameters: lower and upper Miller indices (hkl) of the slabs, number of lower and upper slab layers, vacuum thickness of the slabs, minimum and maximum interface area, interface gap, lattice vector (uv) tolerance in percentage, angle tolerance in degrees, and shape filter option to keep only the most square-like interfaces',
    requires=['lower_poscar_path', 'upper_poscar_path', 'lower_hkl', 'upper_hkl'],
    optional=['lower_slab_layers', 'upper_slab_layers', 'slab_vacuum', 'min_area', 'max_area', 'interface_gap', 'uv_tolerance', 'angle_tolerance', 'shape_filter'],
    defaults={
        'lower_slab_layers': 4,
        'upper_slab_layers': 4,
        'slab_vacuum': 15.0,
        'min_area': 50.0,
        'max_area': 500.0,
        'interface_gap': 2.0,
        'uv_tolerance': 5.0,
        'angle_tolerance': 5.0,
        'shape_filter': False,
        },
    prereqs=[],
))
def generate_interface_from_poscars(input: schemas.GenerateInterfaceFromPoscars) -> dict:
    '''
    Generate VASP POSCAR for interface from two given POSCAR files based on specified parameters
    '''
    
    lower_poscar_path = input.lower_poscar_path
    upper_poscar_path = input.upper_poscar_path
    lower_hkl = input.lower_hkl
    upper_hkl = input.upper_hkl
    lower_slab_layers = input.lower_slab_layers
    upper_slab_layers = input.upper_slab_layers
    slab_vacuum = input.slab_vacuum
    min_area = input.min_area
    max_area = input.max_area
    interface_gap = input.interface_gap
    uv_tolerance = input.uv_tolerance
    angle_tolerance = input.angle_tolerance
    shape_filter = input.shape_filter

    try:
        runs_dir = os.environ.get('MASGENT_SESSION_RUNS_DIR')
        
        interfaces_dir = os.path.join(runs_dir, 'interface_maker')
        os.makedirs(interfaces_dir, exist_ok=True)

        run_interface_maker(
            lower_conv=lower_poscar_path,
            upper_conv=upper_poscar_path,
            lower_hkl=lower_hkl,
            upper_hkl=upper_hkl,
            lower_slab_layers=lower_slab_layers,
            upper_slab_layers=upper_slab_layers,
            slab_vacuum=slab_vacuum,
            min_area=min_area,
            max_area=max_area,
            interface_gap=interface_gap,
            uv_tol=uv_tolerance,
            angle_tol=angle_tolerance,
            shape_filter=shape_filter,
            output_dir=interfaces_dir,
        )

        return {
            'status': 'success',
            'message': f'Generated interface POSCAR(s) in {interfaces_dir}.',
            'interface_dir': interfaces_dir,
        }
    
    except Exception as e:
        return {
            'status': 'error',
            'message': f'Interface POSCAR generation failed: {str(e)}'
        }

@with_metadata(schemas.ToolMetadata(
    name='Generate VASP input files and submit bash script for workflow of convergence tests',
    description='Generate VASP workflow of convergence tests for k-points and energy cutoff based on given POSCAR',
    requires=[],
    optional=['poscar_path', 'kpoint_levels', 'encut_levels'],
    defaults={
        'poscar_path': f'{os.environ.get("MASGENT_SESSION_RUNS_DIR")}/POSCAR',
        'kpoint_levels': [1000, 2000, 3000, 4000, 5000],
        'encut_levels': [300, 400, 500, 600, 700],
        },
    prereqs=[],
))
def generate_vasp_workflow_of_convergence_tests(input: schemas.GenerateVaspWorkflowOfConvergenceTests) -> dict:
    '''
    Generate VASP input files and submit bash script for workflow of convergence tests for k-points and energy cutoff based on given POSCAR
    '''
    # color_print(f'\n[Debug: Function Calling] generate_vasp_workflow_of_convergence_tests with input: {input}', 'green')
    
    poscar_path = input.poscar_path
    test_type = input.test_type
    kpoint_levels = input.kpoint_levels
    encut_levels = input.encut_levels

    def test_kpoints():
        for kppa in kpoint_levels:
            vis = MPStaticSet(structure)
            os.makedirs(kpoint_tests_dir, exist_ok=True)
            test_dir = os.path.join(kpoint_tests_dir, f'kppa_{kppa}')
            os.makedirs(test_dir, exist_ok=True)
            vis.incar.write_file(os.path.join(test_dir, 'INCAR'))
            vis.poscar.write_file(os.path.join(test_dir, 'POSCAR'))
            vis.potcar.write_file(os.path.join(test_dir, 'POTCAR'))
            incar_comments = f'# Generated by Masgent for k-point convergence test with kppa = {kppa}.'
            write_comments(os.path.join(test_dir, 'INCAR'), 'incar', incar_comments)
            poscar_comments = f'# Generated by Masgent for k-point convergence test with kppa = {kppa}.'
            write_comments(os.path.join(test_dir, 'POSCAR'), 'poscar', poscar_comments)
            kpoints = Kpoints.automatic_density(structure, kppa=kppa)
            kpoints.write_file(os.path.join(test_dir, 'KPOINTS'))
            kpoint_comments = f'# Generated by Masgent for k-point convergence test with kppa = {kppa}.'
            write_comments(os.path.join(test_dir, 'KPOINTS'), 'kpoints', kpoint_comments)

        script_path = os.path.join(kpoint_tests_dir, 'masgent_submit.sh')
        batch_script = generate_batch_script()
        with open(script_path, 'w') as f:
            f.write(batch_script)

    def test_encut():
        for encut in encut_levels:
            vis = MPStaticSet(structure, user_incar_settings={'ENCUT': encut})
            os.makedirs(encut_tests_dir, exist_ok=True)
            test_dir = os.path.join(encut_tests_dir, f'encut_{encut}')
            os.makedirs(test_dir, exist_ok=True)
            vis.incar.write_file(os.path.join(test_dir, 'INCAR'))
            vis.poscar.write_file(os.path.join(test_dir, 'POSCAR'))
            vis.kpoints.write_file(os.path.join(test_dir, 'KPOINTS'))
            vis.potcar.write_file(os.path.join(test_dir, 'POTCAR'))
            incar_comments = f'# Generated by Masgent for energy cutoff convergence test with ENCUT = {encut}.'
            write_comments(os.path.join(test_dir, 'INCAR'), 'incar', incar_comments)
            poscar_comments = f'# Generated by Masgent for energy cutoff convergence test with ENCUT = {encut}.'
            write_comments(os.path.join(test_dir, 'POSCAR'), 'poscar', poscar_comments)
            kpoint_comments = f'# Generated by Masgent for energy cutoff convergence test with ENCUT = {encut}.'
            write_comments(os.path.join(test_dir, 'KPOINTS'), 'kpoints', kpoint_comments)

        script_path = os.path.join(encut_tests_dir, 'masgent_submit.sh')
        batch_script = generate_batch_script()
        with open(script_path, 'w') as f:
            f.write(batch_script)

    try:
        runs_dir = os.environ.get('MASGENT_SESSION_RUNS_DIR')
        
        convergence_tests_dir = os.path.join(runs_dir, 'convergence_tests')
        kpoint_tests_dir = os.path.join(convergence_tests_dir, 'kpoint_tests')
        encut_tests_dir = os.path.join(convergence_tests_dir, 'encut_tests')

        structure = Structure.from_file(poscar_path)

        if test_type == 'kpoints':
            test_kpoints()
        elif test_type == 'encut':
            test_encut()
        else:
            test_kpoints()
            test_encut()

        kpoint_tests_files = list_files_in_dir(kpoint_tests_dir) if os.path.exists(kpoint_tests_dir) else []
        encut_tests_files = list_files_in_dir(encut_tests_dir) if os.path.exists(encut_tests_dir) else []

        return {
            'status': 'success',
            'message': f'Generated VASP workflow of convergence tests of k-points in {kpoint_tests_dir} and energy cutoff in {encut_tests_dir}.',
            'kpoint_tests_dir': kpoint_tests_dir,
            'encut_tests_dir': encut_tests_dir,
            'kpoint_tests_files': kpoint_tests_files,
            'encut_tests_files': encut_tests_files,
        }
    
    except Exception as e:
        return {
            'status': 'error',
            'message': f'VASP convergence tests workflow generation failed: {str(e)}'
        }

@with_metadata(schemas.ToolMetadata(
    name='Generate VASP input files and submit bash script for workflow of equation of state (EOS) calculations',
    description='Generate VASP input files and submit bash script for workflow of equation of state (EOS) calculations based on given POSCAR',
    requires=[],
    optional=['poscar_path', 'scale_factors'],
    defaults={
        'poscar_path': f'{os.environ.get("MASGENT_SESSION_RUNS_DIR")}/POSCAR',
        'scale_factors': [0.94, 0.96, 0.98, 1.00, 1.02, 1.04, 1.06],
        },
    prereqs=[],
))
def generate_vasp_workflow_of_eos(input: schemas.GenerateVaspWorkflowOfEos) -> dict:
    '''
    Generate VASP input files and submit bash script for workflow of equation of state (EOS) calculations based on given POSCAR
    '''
    # color_print(f'\n[Debug: Function Calling] generate_vasp_workflow_of_eos with input: {input}', 'green')
    
    poscar_path = input.poscar_path
    scale_factors = input.scale_factors

    try:
        runs_dir = os.environ.get('MASGENT_SESSION_RUNS_DIR')
        
        eos_dir = os.path.join(runs_dir, 'eos_calculations')
        os.makedirs(eos_dir, exist_ok=True)
        
        structure = Structure.from_file(poscar_path)

        for scale in scale_factors:
            scaled_structure = structure.copy()
            scaled_structure.scale_lattice(structure.volume * scale)
            vis = MPStaticSet(scaled_structure)
            scale_dir = os.path.join(eos_dir, f'scale_{scale:.3f}')
            os.makedirs(scale_dir, exist_ok=True)
            vis.incar.write_file(os.path.join(scale_dir, 'INCAR'))
            vis.poscar.write_file(os.path.join(scale_dir, 'POSCAR'))
            vis.kpoints.write_file(os.path.join(scale_dir, 'KPOINTS'))
            vis.potcar.write_file(os.path.join(scale_dir, 'POTCAR'))
            incar_comments = f'# Generated by Masgent for EOS calculation with scale factor = {scale:.3f}.'
            write_comments(os.path.join(scale_dir, 'INCAR'), 'incar', incar_comments)
            poscar_comments = f'# Generated by Masgent for EOS calculation with scale factor = {scale:.3f}.'
            write_comments(os.path.join(scale_dir, 'POSCAR'), 'poscar', poscar_comments)
            kpoint_comments = f'# Generated by Masgent for EOS calculation with scale factor = {scale:.3f}.'
            write_comments(os.path.join(scale_dir, 'KPOINTS'), 'kpoints', kpoint_comments)

        script_path = os.path.join(eos_dir, 'masgent_submit.sh')
        batch_script = generate_batch_script()
        with open(script_path, 'w') as f:
            f.write(batch_script)
        
        eos_files = list_files_in_dir(eos_dir) if os.path.exists(eos_dir) else []

        return {
            'status': 'success',
            'message': f'Generated VASP workflow of EOS calculations in {eos_dir}.',
            'eos_dir': eos_dir,
            'eos_files': eos_files,
        }
    
    except Exception as e:
        return {
            'status': 'error',
            'message': f'VASP EOS workflow generation failed: {str(e)}'
        }

@with_metadata(schemas.ToolMetadata(
    name='Generate VASP input files and submit bash script for workflow of elastic constants calculations',
    description='Generate VASP input files and submit bash script for workflow of elastic constants calculations based on given POSCAR',
    requires=[],
    optional=['poscar_path'],
    defaults={'poscar_path': f'{os.environ.get("MASGENT_SESSION_RUNS_DIR")}/POSCAR'},
    prereqs=[],
))
def generate_vasp_workflow_of_elastic_constants(input: schemas.GenerateVaspWorkflowOfElasticConstants) -> dict:
    '''
    Generate VASP input files and submit bash script for workflow of elastic constants calculations based on given POSCAR
    '''
    # color_print(f'\n[Debug: Function Calling] generate_vasp_workflow_of_elastic_constants with input: {input}', 'green')
    
    poscar_path = input.poscar_path

    try:
        runs_dir = os.environ.get('MASGENT_SESSION_RUNS_DIR')
        
        elastic_dir = os.path.join(runs_dir, 'elastic_constants')
        os.makedirs(elastic_dir, exist_ok=True)
        
        structure = Structure.from_file(poscar_path)
        lattice_matrix = structure.lattice.matrix

        # Create deformation matrixes
        xx = [-0.010, 0.010]
        yy = [-0.010, 0.010]
        zz = [-0.010, 0.010]
        xy = [-0.005, 0.005]
        yz = [-0.005, 0.005]
        xz = [-0.005, 0.005]
        D_xyz = {f'00_strain_xyz': [[1.000, 0.000, 0.000], [0.000, 1.000, 0.000], [0.000, 0.000, 1.000]]}
        D_xx0 = {f'01_strain_xx_{float(xx[0]):.3f}': [[1.000+xx[0], 0.000, 0.000], [0.000, 1.000, 0.000], [0.000, 0.000, 1.000]]}
        D_yy0 = {f'03_strain_yy_{float(yy[0]):.3f}': [[1.000, 0.000, 0.000], [0.000, 1.000+yy[0], 0.000], [0.000, 0.000, 1.000]]}
        D_zz0 = {f'05_strain_zz_{float(zz[0]):.3f}': [[1.000, 0.000, 0.000], [0.000, 1.000, 0.000], [0.000, 0.000, 1.000+zz[0]]]}
        D_xy0 = {f'07_strain_xy_{float(xy[0]):.3f}': [[1.000, xy[0], 0.000], [xy[0], 1.000, 0.000], [0.000, 0.000, 1.000]]}
        D_yz0 = {f'09_strain_yz_{float(yz[0]):.3f}': [[1.000, 0.000, 0.000], [0.000, 1.000, yz[0]], [0.000, yz[0], 1.000]]}
        D_xz0 = {f'11_strain_xz_{float(xz[0]):.3f}': [[1.000, 0.000, xz[0]], [0.000, 1.000, 0.000], [xz[0], 0.000, 1.000]]}
        D_xx1 = {f'02_strain_xx_{float(xx[1]):.3f}': [[1.000+xx[1], 0.000, 0.000], [0.000, 1.000, 0.000], [0.000, 0.000, 1.000]]}
        D_yy1 = {f'04_strain_yy_{float(yy[1]):.3f}': [[1.000, 0.000, 0.000], [0.000, 1.000+yy[1], 0.000], [0.000, 0.000, 1.000]]}
        D_zz1 = {f'06_strain_zz_{float(zz[1]):.3f}': [[1.000, 0.000, 0.000], [0.000, 1.000, 0.000], [0.000, 0.000, 1.000+zz[1]]]}
        D_xy1 = {f'08_strain_xy_{float(xy[1]):.3f}': [[1.000, xy[1], 0.000], [xy[1], 1.000, 0.000], [0.000, 0.000, 1.000]]}
        D_yz1 = {f'10_strain_yz_{float(yz[1]):.3f}': [[1.000, 0.000, 0.000], [0.000, 1.000, yz[1]], [0.000, yz[1], 1.000]]}
        D_xz1 = {f'12_strain_xz_{float(xz[1]):.3f}': [[1.000, 0.000, xz[1]], [0.000, 1.000, 0.000], [xz[1], 0.000, 1.000]]}
        D_all = [D_xyz, D_xx0, D_yy0, D_zz0, D_xy0, D_yz0, D_xz0, D_xx1, D_yy1, D_zz1, D_xy1, D_yz1, D_xz1]

        for D in D_all:
            for key, value in D.items():
                deformation_matrix = np.array(value)
                lattice_matrix_deformed = np.dot(lattice_matrix, deformation_matrix)
                structure_deformed = structure.copy()
                structure_deformed = Structure(lattice=lattice_matrix_deformed, species=structure.species, coords=structure.frac_coords)

                vis = MVLElasticSet(structure_deformed)
                deform_dir = os.path.join(elastic_dir, key)
                os.makedirs(deform_dir, exist_ok=True)
                vis.incar.write_file(os.path.join(deform_dir, 'INCAR'))
                vis.poscar.write_file(os.path.join(deform_dir, 'POSCAR'))
                vis.kpoints.write_file(os.path.join(deform_dir, 'KPOINTS'))
                vis.potcar.write_file(os.path.join(deform_dir, 'POTCAR'))
                incar_comments = f'# Generated by Masgent for elastic constants calculation with deformation {key}.'
                write_comments(os.path.join(deform_dir, 'INCAR'), 'incar', incar_comments)
                poscar_comments = f'# Generated by Masgent for elastic constants calculation with deformation {key}.'
                write_comments(os.path.join(deform_dir, 'POSCAR'), 'poscar', poscar_comments)
                kpoint_comments = f'# Generated by Masgent for elastic constants calculation with deformation {key}.'
                write_comments(os.path.join(deform_dir, 'KPOINTS'), 'kpoints', kpoint_comments)
        
        script_path = os.path.join(elastic_dir, 'masgent_submit.sh')
        batch_script = generate_batch_script()
        with open(script_path, 'w') as f:
            f.write(batch_script)

        elastic_files = list_files_in_dir(elastic_dir) if os.path.exists(elastic_dir) else []

        return {
            'status': 'success',
            'message': f'Generated VASP workflow of elastic constants calculations in {elastic_dir}.',
            'elastic_dir': elastic_dir,
            'elastic_files': elastic_files,
        }

    except Exception as e:
        return {
            'status': 'error',
            'message': f'VASP elastic constants workflow generation failed: {str(e)}'
        }

