#!/usr/bin/env python
# -*- encoding: utf-8 -*-
'''
@file               :ppi.py
@Description:       :
@Date               :2023/09/01 16:56:22
@Author             :lyzeng
@Email              :pylyzeng@gmail.com
@version            :1.0
'''
import sys
from dataclasses import dataclass, field, InitVar
from functools import partial
from pymol import cmd
from pymol.cmd import util
from pathlib import Path
from collections import defaultdict

# local import 
from .pymolbase import PyMOLBase

@dataclass
class PyMOLConcrete(PyMOLBase):
    file: Path
    ligand_positions: defaultdict = field(default_factory=lambda: defaultdict(list), init=False)
    
    def __post_init__(self, file):
        self.load_file(file)
        self.global_obj_name = self.cleanpdb(file.stem)

    def load_file(self, file):
        # 加载文件到PyMOL（这里是一个示例，您可能需要使用实际的PyMOL命令）
        cmd.load(file)
        pass

    def collect_ligand_positions(self, model, chain, resn, resi):
        position = f'/{model}//{chain}/{resi}'
        self.ligand_positions[resn].append(position)

    def gather_ligand_info(self):
        cmd.iterate('organ', 'collect_ligand_positions(model, chain, resn, resi)', space={'collect_ligand_positions': partial(self.collect_ligand_positions)})


    @staticmethod
    def cleanpdb(before_string):
        """
        处理传入参数pdb 例如: pdb1b12.ent 1b12.pdb 转化成 1b12
        """
        # print(f'try to format {before_string}')
        if '.ent' in before_string:
            before_string = before_string[3:-4].lower() 
        elif '.pdb' in before_string:
            before_string = before_string[:4].lower()
        elif len(before_string) == 4:
            before_string = before_string.lower()
        else:
            if len(before_string) != 4:
                raise ValueError(f'length out of 4 {before_string}')
        return before_string

@dataclass
class PDBAnalyzer:
    pdb: Path  # 使用Path对象代替字符串
    ChA: str
    ChB: str
    dist: str
    
    
    # 主运行函数
    def run(self):
        cmd.delete('all')
        cmd.load(str(self.pdb))  # 将Path对象转为字符串以兼容PyMOL命令
        pdb_name = self.pdb.stem  # 获取不包括扩展名的文件名
        
        # ...（省略了其他设置和选择）

        # 调用封装好的函数
        self.HBond()
        self.Pi_Pi_interaction()
        self.Cation_Pi_interaction()
        self.Salt_bridge_positive()
        self.Salt_bridge_negative()

        # ...（省略了其他设置和保存）
        
        cmd.save(f'{pdb_name}_interaction.pse')

class 

# 主程序
def main():
    pdb_file = Path(sys.argv[1])  # 使用Path对象
    CHA = sys.argv[2]
    CHB = sys.argv[3]
    dist = sys.argv[4]
    cmd.set('bg_rgb', 'white')
    
    analyzer = PDBAnalyzer(pdb_file, CHA, CHB, dist)
    analyzer.run()

if __name__ == "__main__":
    main()