##
import pandas as pd
import re

class StructureDataFile:

    def __init__(self):

        self.molecules = []
        self.molecule_number = 0

    #还需要改
    def input(self, filepath):

        try:
            f = open(filepath,'r')
            self.txt = f.read()
        except:
            f = open(filepath, 'r', encoding='ISO-8859-1')
            self.txt = f.read()

        molecules = self.txt.split('$$$$\n')
        self.molecule_number = len(molecules)
        for molecule in molecules:
            information = molecule.split('END\n')
            if len(information) != 2:
                continue
            new_molecule = [information[0]+'END\n']
            new_molecule.append(self._get_keys_and_values(information[1]))
            self.molecules.append(new_molecule)

    def _get_keys_and_values(self, content):
        content_new = content.split('\n\n')
        dic = {}
        for item in content_new[:-1]:
            key = re.findall(r'<([^<>]+)>',item)[0]
            if len(key) > 0:
                value = item.split('\n')[1]
                dic.update({key:value})
            else:
                continue

        return dic

    #还需要改
    def output_molecule(self, molecule_index, filepath=None, property_index=1):

        content = self.molecules[molecule_index][0]
        for key, value in self.molecules[molecule_index][1].items():
            temp = f'> <{key}> ({property_index})\n{value}\n\n'
            content += temp
        if filepath:
            with open(filepath,'w') as f:
                f.write(content)
        else:
            return content

    @property
    def length(self):

        return len(self.molecules)

    def output_molecules(self, molecule_index_list=None, filepath=None):

        if not molecule_index_list:
            molecule_index_list = [i for i in range(self.length)]
        content = ''
        for i in molecule_index_list:
            content += self.output_molecule(i, property_index=i)
            content += '$$$$\n'
        if filepath:
            with open(filepath, 'w') as f:
                f.write(content)
        else:
            return content

    #抓取等于条件的分子级数据，输入为一个字典（key为属性，value为值），后期需要改成支持动态代，返回index
    def grep_molecule_data(self, **kwargs):

        result = []
        for i,molecule in enumerate(self.molecules):
            for key, value in kwargs.items():
                if key in molecule[1].keys():
                    molecule_value = molecule[1][key]
                    if molecule_value == value:
                        result.append(i)
        return result

    #设置key值，没有会自动创建
    def set_molecule_data(self, molecule_index, **kwargs):

        for key, value in kwargs.items():
            self.molecules[molecule_index][1][key] = value

class GlideCovalentDocking:

    def __init__(self,glide_filepath, sdf_filepath):

        self.input_sdf(sdf_filepath)
        self.input_glide_txt(glide_filepath)
        #self.df_clean = self.make_clean_data_copy()

    def input_glide_txt(self, filepath):

        self.df = pd.read_csv(filepath)

    #逻辑有问题，需要修改
    def make_clean_data_copy(self):

        df = self.df
        for column_name in df.columns:
            if '\"' in df[column_name][0]:
                df[column_name] = df[column_name].str.strip('\"')
            else:
                df[column_name] = df[column_name].astype(float)
        self.df_clean = df

    # 抓取等于条件的数据，输入为一个字典（key为属性，value为值），后期需要改成支持动态代码
    def grep_glide_information(self, **kwargs):

        df = self.df_clean
        result_index = []
        for key, value in kwargs.items():
            result_index.extend(df.index[df[key]==value])
        result_index = list(set(result_index))

        return result_index

    def output_glide_txt(self, filepath, df_index=None):

        if not df_index:
            df = self.df
        else:
            df = self.df.iloc[df_index]

        df.to_csv(filepath)

    def input_sdf(self,filepath):

        self.sdf = StructureDataFile()
        self.sdf.input(filepath)

    def merge_glide_txt_and_sdf(self, reaction_type, filepath='result.sdf'):

        content = ''
        for index, row in self.df.iterrows():
            sdf_index = self.sdf.grep_molecule_data(MOLNAME=row['MOLNAME'])[0]
            self.sdf.set_molecule_data(sdf_index, DockingScore=row['docking score'])
            self.sdf.set_molecule_data(sdf_index, ReactionType=reaction_type)
            content += self.sdf.output_molecule(molecule_index=sdf_index,property_index=index+1)
            content += '$$$$\n'
        with open(filepath, 'w') as f:
            f.write(content)

if __name__ == '__main__':

    import argparse

    paser = argparse.ArgumentParser()
    paser.add_argument('-g', help='glide text or csv file')
    paser.add_argument('-s', help='sdf file with MOLNAME')
    paser.add_argument('-o', help='output filepath', default='result.sdf')
    paser.add_argument('-r', help='reaction type', default='NucleophlicDouble')
    args = paser.parse_args()
    mygcd = GlideCovalentDocking(glide_filepath=args.g,sdf_filepath=args.s)
    mygcd.merge_glide_txt_and_sdf(reaction_type=args.r, filepath=args.o)