#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""蛋白数据层构建：UniProt 序列/结构域 + PDB 结构映射
首批：Taq / MMLV RT / UNG / Tth
输出：数据/蛋白结构/<acc>.json + <pdb>.pdb
"""
import io
import json
import os
import urllib.parse
import urllib.request

W = '/Users/liuxinyuan/Desktop/Hermes输出-工作类'
OUT = W + '/数据/蛋白结构'
os.makedirs(OUT, exist_ok=True)


def get_json(u, t=60):
    req = urllib.request.Request(u, headers={'User-Agent': 'biori-protein/1.0', 'Accept': 'application/json'})
    with urllib.request.urlopen(req, timeout=t) as r:
        return json.load(r)


def get_bytes(u, t=120):
    req = urllib.request.Request(u, headers={'User-Agent': 'biori-protein/1.0'})
    with urllib.request.urlopen(req, timeout=t) as r:
        return r.read()


def uniprot(acc):
    return get_json('https://rest.uniprot.org/uniprotkb/%s.json' % acc)


def find_ung():
    """E. coli / T. thermophilus / Gadus morhua 的 UNG"""
    out = []
    for q in ('gene:ung AND organism_id:562', 'gene:ung AND organism_id:274',
              'protein_name:"uracil-DNA glycosylase" AND organism_name:"Gadus morhua"'):
        d = get_json('https://rest.uniprot.org/uniprotkb/search?format=json&size=5&query=' + urllib.parse.quote(q))
        for x in d.get('results', [])[:3]:
            out.append((q.split('AND')[1].strip(), x['primaryAccession'], x['sequence']['length'],
                        x.get('proteinDescription', {}).get('recommendedName', {}).get('fullName', {}).get('value', '')))
    return out


TARGETS = [('Taq', 'P19821'), ('MMLV', 'P03355'), ('Tth', 'P52028'), ('UNG', 'P12295')]
ung = []
print('UNG 候选：')
for u in ung:
    print('   ', u)

# 选 E. coli UNG（gene ung, 562）
ung_acc = None
for cond, acc, ln, name in ung:
    if '562' in cond and 'glycosylase' in (name or '').lower():
        ung_acc = acc
        break
if not ung_acc:
    for cond, acc, ln, name in ung:
        if '562' in cond:
            ung_acc = acc
            break
print('选定 UNG:', ung_acc)
if ung_acc:
    TARGETS.append(('UNG', ung_acc))

summary = []
for label, acc in TARGETS:
    d = uniprot(acc)
    seq = d['sequence']['value']
    feats = [f for f in d.get('features', [])
             if f['type'] in ('Domain', 'Region', 'Active site', 'Metal binding', 'Site', 'Motif', 'Binding site')]
    pdbs = []
    for x in d.get('uniProtKBCrossReferences', []):
        if x.get('database') == 'PDB':
            props = {p['key']: p['value'] for p in x.get('properties', [])}
            pdbs.append({'id': x['id'], 'method': props.get('Method', ''),
                         'resolution': props.get('Resolution', ''), 'chains': props.get('Chains', '')})

    def rank(p):
        r = p['resolution'].replace('A', '').strip()
        try:
            return float(r)
        except Exception:
            return 99
    pdbs.sort(key=rank)
    best = pdbs[0] if pdbs else None
    rec = {
        'label': label, 'accession': acc,
        'name': d.get('proteinDescription', {}).get('recommendedName', {}).get('fullName', {}).get('value', ''),
        'organism': d.get('organism', {}).get('scientificName', ''),
        'length': len(seq), 'sequence': seq,
        'features': [{'type': f['type'], 'start': f['location']['start']['value'],
                      'end': f['location']['end']['value'],
                      'desc': f.get('description', '')} for f in feats],
        'pdb': pdbs[:8], 'best_pdb': best['id'] if best else '',
    }
    io.open('%s/%s.json' % (OUT, acc), 'w', encoding='utf-8').write(json.dumps(rec, ensure_ascii=False, indent=1))
    if best:
        pdb_path = '%s/%s.pdb' % (OUT, best['id'])
        if not os.path.exists(pdb_path):
            data = get_bytes('https://files.rcsb.org/download/%s.pdb' % best['id'])
            io.open(pdb_path, 'wb').write(data)
    summary.append(rec)
    print('\n█ %-5s %s ｜%d aa ｜%s' % (label, acc, len(seq), rec['name'][:50]))
    print('   结构域/位点 %d 条：' % len(feats))
    for f in rec['features'][:6]:
        print('     %-12s %4d-%-4d %s' % (f['type'], f['start'], f['end'], f['desc'][:50]))
    print('   PDB %d 个，最优 %s（%s Å, %s）' % (len(pdbs), rec['best_pdb'], best['resolution'] if best else '-', best['method'] if best else ''))

io.open(OUT + '/_index.json', 'w', encoding='utf-8').write(json.dumps(
    [{'label': r['label'], 'accession': r['accession'], 'name': r['name'], 'organism': r['organism'],
      'length': r['length'], 'best_pdb': r['best_pdb'], 'n_features': len(r['features'])} for r in summary],
    ensure_ascii=False, indent=1))
print('\n✅ 已写入', OUT)