#!/usr/bin/env python
# -*- encoding: utf-8 -*-
'''
@file               :method.py
@Description:       :pymol中一些常用到的一些方法
@Date               :2023/09/24 13:28:30
@Author             :lyzeng
@Email              :pylyzeng@gmail.com
@version            :1.0
'''
import sys
from pathlib import Path
here = Path(__file__).absolute().parent
sys.path.append(here.as_posix())
from plugin.cmd_singleton import cmd
from pymol.preset import ion_sele,lig_sele,solv_sele

def color_mole(selection='hetatm'):
    cmd.color('color9',f'{selection} in name c*')
    cmd.color('color6',f'{selection} in name o*')
    cmd.color('color5',f'{selection} in name s*')
    cmd.color('color3',f'{selection} in name n*')
    cmd.color('color8',f'{selection} in name cl*')
    cmd.color('color10',f'{selection} in name f*')

def focus_ligand(selection: str = 'sele'):
    '''
    watch ligand and protein structure
    '''
    cmd.remove(solv_sele) # 移除金属离子和溶剂
    cmd.remove("resn SO4 or resn PO4 or resn CL")
    cmd.pretty(selection)

def grey(selection='sele',opt=True):
    if opt: optimisation(selection) # 优化
    cmd.color('grey1',f'{selection}')# 对某个对象进行灰化

def optimisation(selection, grey=True):
    if grey: cmd.set('cartoon_color','grey1', selection=selection)
    cmd.set('stick_transparency','0.1', selection=selection)
    cmd.set("ray_shadows","off", selection=selection)
    cmd.set('cartoon_highlight_color', selection=selection)
    cmd.set('cartoon_fancy_helices', selection=selection)
    cmd.set('cartoon_transparency','0.5', selection=selection)