#!/usr/bin/env python
# coding: utf-8

# # Pose & PDBinfo Layer
# 
# @Author: 吴炜坤
# 
# @email：weikun.wu@xtalpi.com/weikunwu@163.com
# 
# PDBinfo是Pose和PDB中信息交换和储存的重要媒介。Pose通常是从PDB文件中衍生出来的，除了原子的坐标信息以外，PDB文件中包含了许多额外的信息，而这些信息是储存在PDBinfo中。比如温度因子数据(bfactor)、晶体解析数据(crystinfo)、原子的占用率(occupancy)等。使用PDBinfo可以实现Pose编号与PDB编号的转换以及Pose的序列信息获取等功能。如果Pose中的氨基酸发生了插入和删除，又或者其他的PDB相关信息发生了变化，更新的信息就需要从当前的Pose中获取并转换成PDBinfo的数据，因此PDBinfo和Pose是实时相互连通的两个信息储存器。

# ### 一、PDB编号与Pose编号

# 在PDB文件中，每个氨基酸都有自己独立的编号，并且氨基酸的PDB编号是依赖于其所在的链ID，如1A，2A...120A，1B，2B...130B等，分别代表A链和B链上的氨基酸位置。
# 而在Pose的概念中，氨基酸的编号是忽略链的分隔，按照PDB文件中链编写的顺序，**从1开始递增**，由于缺乏直观的对应方式，在Pose中的氨基酸和PDB编号的处理往往是相当棘手的，但是现在可以通过PDB_info这个的功能来轻松解决编号转换的问题。

# In[1]:


# 首先依然是从PDB中读入Pose
from pyrosetta import init, pose_from_pdb
init()
pose = pose_from_pdb('./data/4R80.clean.pdb')

# 获取PDBinfo对象:
pose_pdbinfo = pose.pdb_info()


# **编号转换的处理主要使用的是pdb_info中的pdb2pose和pose2pdb函数。**
# - pdb2pose: 将pdb编号翻译成pose编号;
# - pose2pdb: 将pose编号转换成pdb编号;

# In[2]:


# 获取PDB号为24A的氨基酸残基所在的Pose残基编号
pose_number = pose_pdbinfo.pdb2pose(chain='A', res=24)
print(pose_number)


# In[3]:


# 获取Pose残基编号为24的氨基酸残基所在的PDB号
pdb_number = pose_pdbinfo.pose2pdb(res=24)
print(pdb_number)


# ### 二、PDB链信息
# 由于Pose中链号也是从1号开始编，不同于PDB文件中链号信息以字母的形式进行储存。链中信息获取的方式有多种途径:
# - 根据链的PDB chain ID获取Pose chain ID;
# - 根据链的chain ID获取PDB chain ID;
# - 根据氨基酸的Pose编号获取其所在的chain ID;
# - Pose的链氨基酸序列信息;
# - 获取链起始和末端氨基酸信息;

# In[4]:


# 先来看看pdbinfo中链的基本信息:
pose_pdbinfo.short_desc()


# 可见我们的这个pose中共有2条链，分别为A和B链，分别都有76个氨基酸

# In[5]:


# 根据链的PDB chain ID获取Pose chain ID;
from pyrosetta.rosetta.core.pose import get_chain_id_from_chain
chainA_pose_chain_id = get_chain_id_from_chain('A', pose)
chainB_pose_chain_id = get_chain_id_from_chain('B', pose)
print(chainA_pose_chain_id, chainB_pose_chain_id)


# In[6]:


from pyrosetta.rosetta.core.pose import get_chain_from_chain_id
chain1_pdb_chain_id = get_chain_from_chain_id(1, pose)
chain2_pdb_chain_id = get_chain_from_chain_id(2, pose)
print(chain1_pdb_chain_id, chain2_pdb_chain_id)


# In[7]:


# 根据某个氨基酸残基的Pose编号获取其所在的PDB链ID;
residue1_chain_id = pose_pdbinfo.chain(10)
residue82_chain_id = pose_pdbinfo.chain(82)
print(residue1_chain_id, residue82_chain_id)


# In[8]:


# 获取链的的起始和结尾的氨基酸Pose编号:
chain1_start_pose_id = pose.chain_begin(1) # 返回某链起始的rosetta index
chain1_end_pose_id = pose.chain_end(1) # 返回某链终止的rosetta index
print(chain1_start_pose_id, chain1_end_pose_id)


# In[9]:


# 获取链的序列信息:
pose.chain_sequence(1) # 返回某链的氨基酸序列


# ### 三、PDB中晶体解析信息的提取
# 除了基本的编号信息以外，一些晶体相关的信息也可以轻松进行提取:

# In[10]:


# 提取第一个原子的bfactor信息
pose.pdb_info().bfactor(res=1, atom_index=1)  # 返回温度因子信息


# In[11]:


# 获取PDB的晶体信息
crystinfo = pose.pdb_info().crystinfo()
print(crystinfo)


# In[12]:


# 获取原子的occupancy
occupancy = pose.pdb_info().occupancy(res=1, atom_index=1)
print(occupancy)


# ### 练习
# 请写一个小程序，对残基中每一个原子的bfactor进行加和平均处理。

# ### 进阶技巧. 氨基酸残基PDB LABEL
# PyRosetta中的PDBinfo除了储存一些实验信息以外，还可以储存用户的自定义信息，比如用户可以通过pose的label系统对一些氨基酸打上PDB标签，在后续的氨基酸范围选取中快捷方便的操作。
# 进阶部分主要介绍:
# - add_reslabel/get_reslabels/clear_reslabel/res_haslabel
# - icode/set_resinfo

# #### 1. 通过pdbinfo的add_reslabel函数对残基打标签:
# 

# In[13]:


# 打标签
pose_pdbinfo.add_reslabel(1, 'starts')
pose_pdbinfo.add_reslabel(2, 'haha')
pose_pdbinfo.add_reslabel(3, 'end')


# In[14]:


# 查标签
print(pose_pdbinfo.get_reslabels(1))
print(pose_pdbinfo.get_reslabels(2))
print(pose_pdbinfo.get_reslabels(3))
pose.dump_pdb('./data/reslabeled.pdb')


# 在输出PDB文件之后，在文件的最下方可以看到以下的信息被记录:
# - REMARK PDBinfo-LABEL:    1 starts
# - REMARK PDBinfo-LABEL:    2 haha
# - REMARK PDBinfo-LABEL:    3 end

# In[15]:


# 判断标签
pose_pdbinfo.res_haslabel(res=1, target_label='haha')


# In[16]:


# 清除标签
pose_pdbinfo.clear_reslabel(1)
pose_pdbinfo.clear_reslabel(2)
pose_pdbinfo.clear_reslabel(3)


# In[17]:


# 判断标签还是否存在？
pose_pdbinfo.res_haslabel(res=1, target_label='starts')


# #### 2. 关于icode插入与管理
# 在PDB文件中，比如处理抗体结构时，一些特殊的PDB编号(含有insert code）, 如 1A 2B 3C等(此处的代码并非是链号)。
# 通过PDBinfo，用户也可以很方便的插入这些字符。

# In[18]:


# set_resinfo(res: int, chain_id: str, pdb_res: int, ins_code: str) -> None
pose_pdbinfo.set_resinfo(res=1, chain_id='A', pdb_res=1, ins_code='m')


# In[19]:


# 查询某个残基的icode:
print(pose_pdbinfo.icode(res=1))

# 打印1号氨基酸残基信息:
print(pose.residue(1))


# In[20]:


# 保存PDB
pose.dump_pdb('./data/insert_icode.pdb')


# 使用文本编辑器打开后可以见，第一个氨基酸的残基编号发生了变化:从1变成了1m
# ```
# ATOM      1  N   PRO A   1m     35.432  -0.708   7.647  1.00 49.13           N  
# ATOM      2  CA  PRO A   1m     35.959   0.478   8.332  1.00 38.65           C  
# ATOM      3  C   PRO A   1m     36.620   1.469   7.374  1.00 28.53           C  
# ATOM      4  O   PRO A   1m     36.946   1.110   6.240  1.00 28.02           O  
# ATOM      5  CB  PRO A   1m     36.987  -0.100   9.317  1.00 32.04           C  
# ATOM      6  CG  PRO A   1m     37.119  -1.563   8.970  1.00 40.69           C  
# ATOM      7  CD  PRO A   1m     35.846  -1.957   8.305  1.00 42.21           C  
# ATOM      8  HA  PRO A   1m     35.141   0.977   8.872  1.00  0.00           H  
# ATOM      9 1HB  PRO A   1m     37.944   0.434   9.218  1.00  0.00           H  
# ATOM     10 2HB  PRO A   1m     36.642   0.048  10.351  1.00  0.00           H  
# ATOM     11 1HG  PRO A   1m     37.985  -1.720   8.311  1.00  0.00           H  
# ATOM     12 2HG  PRO A   1m     37.300  -2.154   9.880  1.00  0.00           H  
# ATOM     13 1HD  PRO A   1m     36.044  -2.759   7.579  1.00  0.00           H  
# ATOM     14 2HD  PRO A   1m     35.123  -2.291   9.064  1.00  0.00           H  
# ATOM     15 1H   PRO A   1m     35.759  -0.700   6.700  1.00  0.00           H  
# ATOM     16 2H   PRO A   1m     34.427  -0.649   7.643  1.00  0.00           H
# ```

# In[ ]: