from rdkit.Chem import AllChem
from pubchemfp import GetPubChemFPs
from rdkit import Chem
import paddle
import paddle.nn as nn

class FPN(nn.Layer):
    def __init__(self, dropout, cuda, hidden_size):
        super(FPN, self).__init__()
        self.dropout= dropout
        self.cuda = cuda
        self.hidden_dim = hidden_size
        self.fc = nn.Sequential(
            nn.Dropout(self.dropout),
            nn.Linear(in_features=2513, out_features=2000, weight_attr=nn.initializer.KaimingNormal()),
            nn.Swish(),
            nn.Dropout(self.dropout),
            nn.Linear(in_features=2000, out_features=1000, weight_attr=nn.initializer.KaimingNormal()),
            nn.Swish(),
            nn.Dropout(self.dropout),
            nn.Linear(in_features=1000, out_features=self.hidden_dim, weight_attr=nn.initializer.KaimingNormal())
        )

    def forward(self, smile):
        fp_list = []
        for i, one in enumerate(smile):
            fp = []
            mol = Chem.MolFromSmiles(one)  # 从smiles编码中获取结构信息

            fp_maccs = AllChem.GetMACCSKeysFingerprint(mol)  # 获得MACC指纹信息
            fp_phaErGfp = AllChem.GetErGFingerprint(mol, fuzzIncrement=0.3, maxPath=21, minPath=1)  # 获得ERG指纹信息
            fp_pubcfp = GetPubChemFPs(mol)  # 获得PubChem指纹信息
            fp_morgan = AllChem.GetMorganFingerprintAsBitVect(mol, 2, nBits=1024)
            fp.extend(fp_maccs)
            fp.extend(fp_phaErGfp)
            fp.extend(fp_pubcfp)
            fp.extend(fp_morgan)
            fp_list.append(fp)

        fp_list = paddle.to_tensor(fp_list,dtype='float32')
        if self.cuda:
            fp_list = fp_list.cuda()

        fpn_out = self.fc(fp_list)  # 把指纹信息放入全连接网络

        return fpn_out