#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""MD 轨迹 → 动画（MP4 + GIF）
用法：python md_movie.py <md_dir> [帧数]
依赖：mdtraj（md 环境）、pymol（pymol 环境）、ffmpeg
"""
import os
import shutil
import subprocess
import sys

import mdtraj as md

d = sys.argv[1]
n = int(sys.argv[2]) if len(sys.argv) > 2 else 120
W = '/Users/liuxinyuan/Desktop/Hermes输出-工作类'
MAMBA = os.path.expanduser('~/.local/bin/micromamba')

import tempfile
_tmp = tempfile.mkdtemp(prefix='md_')
shutil.copy(os.path.join(d, 'final.pdb'), os.path.join(_tmp, 't.pdb'))
shutil.copy(os.path.join(d, 'traj.dcd'), os.path.join(_tmp, 't.dcd'))
traj = md.load(os.path.join(_tmp, 't.dcd'), top=os.path.join(_tmp, 't.pdb'))
# 只渲染主链（残基最多的链），与 RMSD/RMSF 口径一致
chains = list(traj.topology.chains)
main = max(chains, key=lambda c: sum(1 for _ in c.residues))
traj = traj.atom_slice(traj.topology.select('chainid %d' % main.index))
step = max(1, traj.n_frames // n)
sub = traj[::step]
pdb_movie = os.path.join(d, 'movie.pdb')
sub.save_pdb(pdb_movie)
print('轨迹 %d 帧 → 抽 %d 帧' % (traj.n_frames, sub.n_frames))

frames = os.path.join(d, 'frames')
shutil.rmtree(frames, ignore_errors=True)
os.makedirs(frames, exist_ok=True)
pml = os.path.join(d, 'movie.pml')
open(pml, 'w').write('\n'.join([
    'load %s, m' % pdb_movie,
    'hide everything',
    'show cartoon, m',
    'spectrum resi, blue_white_red, m',
    'bg_color white',
    'set ray_opaque_background, off',
    'set antialias, 2',
    'set ray_shadows, 0',
    'select core, resi 1-220',
    'orient core',
    'zoom core, 1.5',
    'mset 1 x%d' % sub.n_frames,
    'mpng %s/, 1, %d' % (frames, sub.n_frames),
]) + '\n')
r = subprocess.run([MAMBA, 'run', '-n', 'pymol', 'pymol', '-cq', pml],
                   capture_output=True, text=True, timeout=1800)
print((r.stdout + r.stderr)[-300:])

mp4 = os.path.join(d, 'md.mp4')
gif = os.path.join(d, 'md.gif')
subprocess.run(['ffmpeg', '-y', '-framerate', '12', '-i', os.path.join(frames, '%04d.png'),
                '-vf', 'scale=720:-1:flags=lanczos', '-pix_fmt', 'yuv420p', mp4],
               capture_output=True)
subprocess.run(['ffmpeg', '-y', '-framerate', '12', '-i', os.path.join(frames, '%04d.png'),
                '-vf', 'scale=560:-1:flags=lanczos,split[a][b];[a]palettegen[p];[b][p]paletteuse',
                gif], capture_output=True)
print('MP4:', os.path.exists(mp4), os.path.getsize(mp4) if os.path.exists(mp4) else 0)
print('GIF:', os.path.exists(gif), os.path.getsize(gif) if os.path.exists(gif) else 0)