#! /usr/bin/env python
# -*- coding: UTF-8 -*-

import os
import pandas as pd
import time

# 读取参数
end_temp = int(input("请输入结束温度(K)："))
step_size = int(input("请输入插点步长(K)："))

# 生成温度列表
temperatures = list(range(step_size, end_temp+1, step_size))

# 创建空的数据框
df = pd.DataFrame(columns=['Temperature', 'Zero-point energy E_ZPE', 'Thermal correction to U(T)', 'Thermal correction to H(T)', 'Thermal correction to G(T)', 'Entropy S(eV/K)', 'Entropy contribution T*S'])

# 循环执行命令并解析输出文件
for temp in temperatures:
    # 执行命令
    os.system(f"(echo 501;echo {temp}) | vaspkit > vaspkit.out")
    
    # 解析输出文件
    with open('vaspkit.out', 'r') as f:
        lines = f.readlines()
        
    # 提取数据并填入数据框
    data = {}
    data['Temperature'] = temp
    for line in lines:
        if line.startswith(' Zero-point energy E_ZPE'):
            data['Zero-point energy E_ZPE'] = float(line.split()[-2])
        elif line.startswith(' Thermal correction to U(T)'):
            data['Thermal correction to U(T)'] = float(line.split()[-2])
        elif line.startswith(' Thermal correction to H(T)'):
            data['Thermal correction to H(T)'] = float(line.split()[-2])
        elif line.startswith(' Thermal correction to G(T)'):
            data['Thermal correction to G(T)'] = float(line.split()[-2])
        elif line.startswith(' Entropy S'):
            data['Entropy S(eV/K)'] = float(line.split()[-2])
        elif line.startswith(' Entropy contribution T*S'):
            data['Entropy contribution T*S'] = float(line.split()[-2])
    df = df.append(data, ignore_index=True)
    
    # 删除vaspkit.out文件
    os.remove('vaspkit.out')
    
    print('计算'+str(temp)+'K中')

# 输出数据框到csv文件
df.to_csv('output.csv', index=False)