#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""通用：AlphaFold 预测 + 实验结构叠加渲染
用法：python3 render_overlay.py <label> <uniprot_acc> <exp_pdb>
"""
import json
import os
import subprocess
import sys
import urllib.request

W = '/Users/liuxinyuan/Desktop/Hermes输出-工作类'
PDB = W + '/数据/蛋白结构'
OUT = W + '/销售分析/蛋白结构图'
MAMBA = os.path.expanduser('~/.local/bin/micromamba')
UA = 'Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 Chrome/126 Safari/537.36'


def get(url, binary=False, t=180):
    req = urllib.request.Request(url, headers={'User-Agent': UA, 'Accept': '*/*'})
    with urllib.request.urlopen(req, timeout=t) as r:
        return r.read() if binary else r.read().decode()


def af_url(acc):
    d = json.loads(get('https://alphafold.ebi.ac.uk/api/prediction/' + acc))
    # 选覆盖最长的片段
    d.sort(key=lambda x: (x.get('uniprotEnd', 0) - x.get('uniprotStart', 0)), reverse=True)
    return d[0]['pdbUrl'], d[0].get('uniprotStart'), d[0].get('uniprotEnd'), d[0].get('globalMetricValue')


def main(label, acc, exp_pdb):
    url, s, e, plddt = af_url(acc)
    af_path = '%s/AF-%s.pdb' % (PDB, acc)
    if not os.path.exists(af_path):
        open(af_path, 'wb').write(get(url, binary=True))
    pml = '\n'.join([
        'load %s, af' % af_path,
        'load %s/%s, exp' % (PDB, exp_pdb),
        'align exp, af',
        'hide everything, exp',
        'show cartoon, af',
        'color grey80, af',
        'color marine, af and byres (af within 2.5 of exp)',
        'bg_color white',
        'set ray_opaque_background, off',
        'set cartoon_fancy_helices, 1',
        'set antialias, 2',
        'zoom af',
        'png %s/%s.png, width=1400, height=1400, dpi=150, ray=1' % (OUT, label),
    ]) + '\n'
    pml_path = '%s/pml/%s_af.pml' % (W + '/脚本', label)
    os.makedirs(os.path.dirname(pml_path), exist_ok=True)
    open(pml_path, 'w').write(pml)
    r = subprocess.run([MAMBA, 'run', '-n', 'pymol', 'pymol', '-cq', pml_path],
                       capture_output=True, text=True, timeout=600)
    for ln in (r.stdout + r.stderr).splitlines():
        if 'RMSD' in ln or 'ScenePNG' in ln or 'Error' in ln:
            print('  ', ln.strip()[:120])
    png = '%s/%s.png' % (OUT, label)
    print('%s ✅ %d bytes ｜AF 区间 %s-%s ｜pLDDT %s' % (label, os.path.getsize(png) if os.path.exists(png) else 0, s, e, plddt))
    return s, e, plddt


if __name__ == '__main__':
    main(sys.argv[1], sys.argv[2], sys.argv[3])