#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""用 PyMOL 渲染四个蛋白的丝带图（按结构域着色）
输出：销售分析/蛋白结构图/<label>.png
"""
import os
import subprocess
from collections import defaultdict

W = '/Users/liuxinyuan/Desktop/Hermes输出-工作类'
PDB = W + '/数据/蛋白结构'
OUT = W + '/销售分析/蛋白结构图'
os.makedirs(OUT, exist_ok=True)
MAMBA = os.path.expanduser('~/.local/bin/micromamba')

# 蛋白 → (PDB, 链, 结构域[(名称, 起, 止, 颜色)])
JOBS = {
    'Taq': ('1TAQ', 'A', [('5-3外切', 1, 290, 'orange'), ('聚合酶', 291, 832, 'marine')]),
    'Tth': ('8XK9_conv', 'A', [('5-3外切', 1, 290, 'orange'), ('聚合酶', 291, 834, 'marine')]),
    'UNG': ('3UF7', 'A', [('UDG 结构域', 1, 229, 'teal')]),
    'MMLV': ('4MH8', 'A', [('RT 结构域', 24, 350, 'marine'), ('RNase H', 351, 487, 'orange')]),
}


def chain_range(pdb, ch):
    lo, hi = 10 ** 6, 0
    fn = '%s/%s' % (PDB, pdb if pdb.endswith('.cif') else pdb + '.pdb')
    for line in open(fn, encoding='utf-8', errors='ignore'):
        if line.startswith(('ATOM', 'HETATM')) and line[12:16].strip() == 'CA' and line[21] == ch:
            r = int(line[22:26])
            lo = min(lo, r); hi = max(hi, r)
    return lo, hi


for label, (pdb, ch, domains) in JOBS.items():
    lo, hi = chain_range(pdb, ch)
    fname = pdb if pdb.endswith('.cif') else pdb + '.pdb'
    pml = ['load %s/%s, m' % (PDB, fname), 'hide everything', 'bg_color white',
           'set ray_opaque_background, off', 'set cartoon_fancy_helices, 1',
           'set antialias, 2', 'set cartoon_transparency, 0.0']
    pml.append('select ch, chain %s and resi %d-%d' % (ch, lo, hi))
    pml.append('show cartoon, ch')
    for idx, (name, a, b, color) in enumerate(domains):
        sel = 'dom%d' % idx
        pml.append('select %s, chain %s and resi %d-%d' % (sel, ch, a, b))
        pml.append('color %s, %s' % (color, sel))
    pml.append('zoom ch')
    pml.append('turn x 12; turn y -18')
    pml.append('png %s/%s.png, width=1400, height=1400, dpi=150, ray=1' % (OUT, label))
    pml_path = '%s/_%s.pml' % (OUT, label)
    open(pml_path, 'w').write('\n'.join(pml) + '\n')
    r = subprocess.run([MAMBA, 'run', '-n', 'pymol', 'pymol', '-cq', pml_path],
                       capture_output=True, text=True, timeout=300)
    ok = os.path.exists('%s/%s.png' % (OUT, label))
    print('%-5s %s ｜chain %s 残基 %d-%d ｜%s' % (
        label, pdb, ch, lo, hi, ('✅ %.0f KB' % (os.path.getsize('%s/%s.png' % (OUT, label)) / 1024)) if ok else '❌ ' + r.stderr[-120:]))

print('\n输出目录:', OUT)