#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""MD 轨迹分析：RMSD / RMSF / Rg（仅主链）
用法：python md_analyze.py <md_dir>
说明：晶体可能含多条链，残基号跨链重排 -> 取残基最多的链为主链，x 轴按顺序编号
"""
import json
import os
import shutil
import sys
import tempfile

import matplotlib
matplotlib.use('Agg')
matplotlib.rcParams['font.family'] = ['PingFang SC', 'Heiti SC', 'Arial Unicode MS', 'DejaVu Sans']
matplotlib.rcParams['axes.unicode_minus'] = False
import matplotlib.pyplot as plt
import mdtraj as md
import numpy as np

d = sys.argv[1]
DT_PS = 5.0  # DCDReporter 每 2500 步 = 5 ps 一帧；DCD 头不存真实时间

_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'))
T_PS = np.arange(traj.n_frames) * DT_PS

chains = list(traj.topology.chains)
main = max(chains, key=lambda c: sum(1 for _ in c.residues))
t = traj.atom_slice(traj.topology.select('chainid %d' % main.index))
print('全体系 %d 链 / %d 残基 / %d 原子 ｜主链 %d 残基 ｜%d 帧 / %.2f ns' % (
    len(chains), traj.n_residues, traj.n_atoms, t.n_residues,
    traj.n_frames, T_PS[-1] / 1000.0))

bb = t.atom_slice(t.topology.select('backbone'))
rmsd = md.rmsd(bb, bb, 0) * 10
ca = t.atom_slice(t.topology.select('name CA'))
rmsf = md.rmsf(ca, ca, 0) * 10
rg = md.compute_rg(t) * 10

res = list(t.topology.residues)
x = np.arange(1, len(res) + 1)

np.savetxt(os.path.join(d, 'rmsd.csv'), np.c_[T_PS / 1000.0, rmsd], delimiter=',',
           header='time_ns,rmsd_A', comments='')
np.savetxt(os.path.join(d, 'rmsf.csv'), np.c_[x, [r.resSeq for r in res], rmsf], delimiter=',',
           header='index,resSeq,rmsf_A', comments='')

fig, ax = plt.subplots(2, 1, figsize=(8, 5.2), dpi=150)
ax[0].plot(T_PS / 1000.0, rmsd, color='#1852B6', lw=1.4)
ax[0].set_xlabel('时间 (ns)'); ax[0].set_ylabel('RMSD (Å)')
ax[0].set_title('主链 RMSD（主链 %d 残基）' % len(res))
ax[0].grid(alpha=.25)
ax[1].plot(x, rmsf, color='#E08A1E', lw=1.2)
ax[1].set_xlabel('残基序号（主链顺序）'); ax[1].set_ylabel('RMSF (Å)')
ax[1].set_title('CA 原子涨落 (RMSF)')
ax[1].grid(alpha=.25)
plt.tight_layout()
plt.savefig(os.path.join(d, 'md_curves.png'))

summary = {
    'frames': int(traj.n_frames), 'atoms': int(t.n_atoms), 'residues': int(t.n_residues),
    'chains_total': len(chains),
    'duration_ns': float(T_PS[-1] / 1000.0),
    'rmsd_mean_A': float(rmsd.mean()), 'rmsd_max_A': float(rmsd.max()),
    'rmsd_first_A': float(rmsd[0]), 'rmsd_last_A': float(rmsd[-1]),
    'rmsf_mean_A': float(rmsf.mean()), 'rmsf_max_A': float(rmsf.max()),
    'rg_mean_A': float(rg.mean()), 'rg_start_A': float(rg[0]), 'rg_end_A': float(rg[-1]),
    'top_flex_residues': [{'index': int(x[i]), 'resSeq': int(res[i].resSeq),
                           'name': str(res[i].name), 'rmsf': round(float(rmsf[i]), 2)}
                          for i in np.argsort(rmsf)[-8:][::-1]],
}
json.dump(summary, open(os.path.join(d, 'summary.json'), 'w'), ensure_ascii=False, indent=1)
print('RMSD 均值 %.2f Å（末帧 %.2f）｜RMSF 均值 %.2f Å ｜Rg %.1f → %.1f Å' % (
    rmsd.mean(), rmsd[-1], rmsf.mean(), rg[0], rg[-1]))
print('柔性最大:', ', '.join('%s%d(%.1fÅ)' % (v['name'], v['resSeq'], v['rmsf'])
                              for v in summary['top_flex_residues'][:5]))
print('输出:', d)