from loguru import logger
from pymol import cmd
import subprocess
from pathlib import Path

# 删除当前目录下的*.mp4和*.png文件
for file_path in Path('.').glob('*.mp4'):
    file_path.unlink()

for file_path in Path('.').glob('*.png'):
    file_path.unlink()

# 设置日志文件
logger.add("pymol_animation.log")
cmd.log_open('pymol_animation_cmd.log')
# 最高渲染质量
cmd.util.performance(0)
cmd.rebuild()
# 初始化并加载4i24结构
cmd.reinitialize('everything')
cmd.fetch('4i24')
cmd.copy('4i24_mutated', '4i24') # 复制4i24到4i24_mutated
# 假设protgrey和colorp.color_mole是你自定义的函数
protgrey('4i24')
protgrey('4i24_mutated')
colorp.color_mole(obj='hetatm')
cmd.hide('everything', '4i24_mutated')
i = '4i24_mutated' # object
j = 'A' # chain
k = '797' # site
mutate_site = f"/{i}//{j}/{k}"
mutation_type = 'GLY'
# 对4i24_mutated进行突变
cmd.select('mut_res', mutate_site)
cmd.wizard("mutagenesis")
cmd.refresh_wizard()
cmd.get_wizard().set_mode(mutation_type)
cmd.get_wizard().do_select(mutate_site)
cmd.get_wizard().apply()
cmd.set_wizard("done")
cmd.show('sticks', mutate_site)
cmd.hide('everything', '4i24_mutated')

# 先打关键帧然后创建视频

# 创建第一个关键帧
selection1 = "/4i24//A/797"
cmd.scene('scene1', 'store')
cmd.orient(selection1)
cmd.show('sticks', selection1)


# 创建第二个关键帧
selection2 = "/4i24//A/1C9"
cmd.select("both_residues", f"{selection1} or {selection2}")
cmd.orient("both_residues")
cmd.scene('scene2', 'store')

# 设置动画和mview
cmd.mset("1 x900")  # 创建一个总共1200帧的动画框架，每秒120帧共计10秒
# 第一个关键帧: 1帧
cmd.mview(action='store', first=1, scene='scene1') 

# 第二个关键帧: 300帧
cmd.mview(action='store', first=300, scene='scene2') 

# 插值以平滑过渡
cmd.mview('interpolate', first=1, last=300)

residue_to_show = "/4i24//A/797"  # 用你想显示的残基替换
cmd.mdo(360, f"show lines, {residue_to_show}; label {residue_to_show}/n, f'res {residue_to_show}(CYS)'")

# 隐藏4i24对象: 600帧
cmd.mdo(600, "hide everything, 4i24")

# 显示4i24_mutated对象: 610帧
cmd.mdo(610, "show cartoon, 4i24_mutated")

# 显示突变残基: 620帧
cmd.mdo(620, "show sticks, /4i24_mutated//A/797")

# 为残基添加标签: 630帧
residue_to_show = "/4i24_mutated//A/797"
cmd.mdo(630, f"show lines, {residue_to_show}; label {residue_to_show}/n, f'res {residue_to_show}(GLY)'")

# 插值以平滑过渡
cmd.mview('interpolate', first=300, last=900)

# 设置viewport尺寸为1080P（1920x1080像素）
cmd.viewport(1920, 1080)
cmd.mplay() # 播放动画

# 光线追踪和保存每一帧
# for frame in range(1, 901):  
#     cmd.frame(frame)
#     logger.info(f"Rendering frame {frame}")  
#     cmd.refresh()
#     cmd.ray(width=1920, height=1080)
#     cmd.png(f"frame_transition_{frame:04d}.png")

# 使用ffmpeg合成视频
# subprocess.run([
#     'ffmpeg', 
#     '-framerate', '120',  # 120帧/秒
#     '-i', 'frame_transition_%04d.png', 
#     '-c:v', 'libx264', 
#     '-pix_fmt', 'yuv420p', 
#     'transition_video.mp4'
# ])


# ffmpeg -framerate 120 -i frame_transition_%04d.png -c:v libx264 -pix_fmt yuv420p transition_video.mp4