#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""AlphaFold 预测 + 实验结构叠加渲染（PyMOL）
用法：python3 render_overlay.py --acc P52028 --exp 8XK9_conv.pdb --label Tth \
        --pdb-dir <结构目录> --out-dir <输出目录>
"""
import argparse
import json
import os
import subprocess
import urllib.request

UA = 'Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 Chrome/126 Safari/537.36'
MAMBA = os.path.expanduser('~/.local/bin/micromamba')


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_best(acc):
    entries = json.loads(get('https://alphafold.ebi.ac.uk/api/prediction/' + acc))
    entries.sort(key=lambda x: x['uniprotEnd'] - x['uniprotStart'], reverse=True)
    e = entries[0]
    return e['pdbUrl'], e['uniprotStart'], e['uniprotEnd'], e.get('globalMetricValue')


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('--acc', required=True)
    ap.add_argument('--exp', required=True, help='实验结构文件名（在 pdb-dir 下）')
    ap.add_argument('--label', required=True)
    ap.add_argument('--pdb-dir', default='data/proteins')
    ap.add_argument('--out-dir', default='out/protein-figures')
    ap.add_argument('--width', type=int, default=1400)
    a = ap.parse_args()

    os.makedirs(a.out_dir, exist_ok=True)
    af_path = os.path.join(a.pdb_dir, 'AF-%s.pdb' % a.acc)
    if not os.path.exists(af_path):
        url, s, e, plddt = af_best(a.acc)
        open(af_path, 'wb').write(get(url, binary=True))
    pml = '\n'.join([
        'load %s, af' % af_path,
        'load %s/%s, exp' % (a.pdb_dir, a.exp),
        '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=%d, height=%d, dpi=150, ray=1' % (a.out_dir, a.label, a.width, a.width),
    ]) + '\n'
    pml_path = os.path.join(a.out_dir, '_%s.pml' % a.label)
    open(pml_path, 'w').write(pml)
    r = subprocess.run([MAMBA, 'run', '-n', 'pymol', 'pymol', '-cq', pml_path],
                       capture_output=True, text=True, timeout=900)
    for ln in (r.stdout + r.stderr).splitlines():
        if 'RMSD' in ln or 'ScenePNG' in ln or 'Error' in ln or 'Parsing' in ln:
            print(ln.strip())
    print('PNG:', os.path.join(a.out_dir, a.label + '.png'))


if __name__ == '__main__':
    main()